xudong963 opened a new issue, #24807:
URL: https://github.com/apache/datafusion/issues/24807
### Describe the bug
`EnsureRequirements` can silently remove a pushed-down `LIMIT` when the
limit is stored as `fetch` on a `SortPreservingMergeExec` or
`CoalescePartitionsExec` and distribution is optimized again.
`remove_dist_changing_operators` records the minimum `fetch` while stripping
distribution-changing operators, but the saved value is only reapplied when
`add_merge_on_top` inserts a replacement operator for a `SinglePartition`
requirement. If no replacement consumes it, `ensure_distribution` returns the
stripped child and drops the saved `fetch`.
For example, current `main` rewrites:
```text
SortPreservingMergeExec: [c@2 ASC], fetch=5
DataSourceExec: ... 2 partitions ... output_ordering=[c@2 ASC]
```
to the multi-partition `DataSourceExec` without any `fetch`. The same
happens to a top-level `CoalescePartitionsExec: fetch=5`.
This silently removes the query's global limit and can return more rows than
requested, so this is a correctness issue rather than only an optimization
difference.
PR #21976 intended to fix the fetch-loss problem from #14150 while
introducing `EnsureRequirements`, but the fallback path where no replacement
distribution operator is needed remains uncovered. This is also related to the
still-open #21169, whose root-cause description refers to the
pre-`EnsureRequirements` implementation.
### To Reproduce
Construct an already optimized, fetched merge over a sorted multi-partition
source and run `EnsureRequirements` again:
```rust
let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]);
let plan: Arc<dyn ExecutionPlan> = Arc::new(
SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5)),
);
let optimized =
EnsureRequirements::new().optimize(plan,
&test_suite_default_config_options())?;
assert_eq!(optimized.fetch(), Some(5));
```
On current `main`, `optimized.fetch()` is `None` and the merge is removed.
Replacing the merge with
`CoalescePartitionsExec::new(input).with_fetch(Some(5))` reproduces the same
loss.
The issue affects DataFusion 55.0.0 and current `main`.
### Expected behavior
Distribution reoptimization must preserve the effective minimum `fetch` from
removed `SortPreservingMergeExec` and `CoalescePartitionsExec` nodes. If no
newly inserted distribution operator consumes it, the optimizer should restore
the original fetch-capable operator around the optimized child.
### Additional context
I have a fix with regression coverage for both fetched operators and for
moving a fetched ordered merge onto a replacement sort. I will open the PR
after filing this issue.
--
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]
