andygrove opened a new issue, #2768:
URL: https://github.com/apache/datafusion-ballista/issues/2768
Disclaimer: this issue was drafted by an LLM (Claude Code) at my request,
against `70a074322`. File/line references and the history were checked against
the code, and the test below was run (it fails on `main` and passes with the
change suggested at the end), but please treat the suggested fix as a
suggestion.
**Describe the bug**
`ballista.planner.propagate_empty.enabled` (default `true`) is documented in
`docs/source/user-guide/configs.md` and exposed in the Python API as
`BallistaSessionConfig.with_propagate_empty_enabled()`, but nothing reads it.
Setting it to `false` does not disable the AQE `PropagateEmptyExecRule`:
- The key, its `ConfigEntry` and `BallistaConfig::propagate_empty_enabled()`
are in `ballista/core/src/config.rs` (lines 172, 468-473 and 1061). The
accessor has no callers.
- `PropagateEmptyExecRule::optimize` ignores its config argument (`_config`,
`propagate_empty.rs:235-241`), and `AdaptivePlanner::default_optimizers_with`
registers the rule unconditionally (`planner.rs:607`).
- #1723 added the key and the accessor but never read them. #2637 added the
Python setter; its test (`test_session_config.py:282-286`) only checks that the
key is written.
**To Reproduce**
This is the flow of `should_propagate_empty_stage`
(`aqe/test/alter_stages.rs`) on a Ballista session config with the flag off.
Stage 0 finishes with no rows, and stage 1 is still collapsed to `EmptyExec`:
```rust
#[tokio::test]
async fn propagate_empty_can_be_disabled() -> datafusion::error::Result<()> {
let config = SessionConfig::new_with_ballista()
.with_target_partitions(2)
.with_round_robin_repartition(false)
.set_bool("ballista.planner.propagate_empty.enabled", false);
let ctx = SessionContext::new_with_state(
SessionStateBuilder::new_with_default_features()
.with_config(config)
.build(),
);
ctx.register_batch("t", mock_batch()?)?;
let plan = ctx
.sql("select min(a) as c0, max(b) as c1, c as c2 from t group by c")
.await?
.create_physical_plan()
.await?;
let mut planner =
AdaptivePlanner::try_from_plan(ctx.state().config(), plan,
"test_job".into())?;
planner.runnable_stages()?.unwrap();
planner.finalise_stage_internal(0,
mock_partitions_with_statistics_no_data())?;
let stages = planner.runnable_stages()?.unwrap();
assert_plan!(stages[0].plan.as_ref(), @ r"
ShuffleWriterExec: partitioning: Hash([c2@2], 2)
ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
AggregateExec: mode=FinalPartitioned, gby=[c@0 as c],
aggr=[min(t.a), max(t.b)]
ShuffleReaderExec: upstream_stage: 0, partitioning: Hash([c@0], 2)
");
Ok(())
}
```
On `main` the stage 1 plan is the same as with the flag on:
```
ShuffleWriterExec: partitioning: RoundRobinBatch(2)
RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
EmptyExec
```
**Expected behavior**
With the flag set to `false`, AQE does not run `PropagateEmptyExecRule`, and
stages downstream of an empty stage run their normal plans (the snapshot above).
**Additional context**
The other AQE rules read their flag inside `optimize` and return the plan
unchanged when it is off (`coalesce_partitions.rs:159-166`,
`optimize_skewed_join.rs:61`, `parallel_window.rs:114`). The same guard here
makes the test above pass, and the rest of the scheduler lib tests still pass:
```rust
let bc = config
.extensions
.get::<BallistaConfig>()
.cloned()
.unwrap_or_default();
if !bc.propagate_empty_enabled() {
return Ok(plan);
}
```
A working off switch is worth having, because this rule has produced wrong
answers before:
- #2185, fixed in #2194: TPC-DS q61 returned no rows instead of one all-NULL
row.
- #2200: grouping sets and embedded projections.
A fix should also run an empty-stage query end to end with the flag off,
since the planner test doesn't execute anything.
The alternative is to remove the key, the accessor, the configs.md row and
the Python setter, but that breaks anyone who already sets it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
