apache / apache/datafusion

Bug: EnforceDistribution optimizer loses fetch (LIMIT) from CoalescePartitionsExec and SortPreservingMergeExec

Open
#21,169 0 comments 0 reactions 1 assignee Claimed by @zhuqi-lucas View on GitHub
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### Describe the bug

When `LimitPushdown` merges a `GlobalLimitExec` into a `CoalescePartitionsExec` (or `SortPreservingMergeExec`) as a `fetch` value, the `EnforceDistribution` optimizer rule strips and re-inserts distribution-changing operators **without preserving the `fetch`**. This causes queries with `LIMIT` over multi-partition sources to silently lose the limit and potentially return duplicate/extra rows.

### Root cause

In `enforce_distribution.rs`, the function `remove_dist_changing_operators` strips `CoalescePartitionsExec`, `SortPreservingMergeExec`, and `RepartitionExec` from the plan tree. It does not capture or propagate any `fetch` value that was embedded in those operators. Later, when `add_merge_on_top` re-inserts a merge operator to satisfy `SinglePartition` distribution, the `fetch` is gone.

### To Reproduce

1. Create a parquet table with multiple row groups / partitions.
2. Run a query with `LIMIT`, e.g. `SELECT * FROM t LIMIT 1`.
3. After `LimitPushdown`, the plan has `CoalescePartitionsExec(fetch=1)`.
4. `EnforceDistribution` strips the `CoalescePartitionsExec` and re-inserts one without `fetch`.
5. The limit is silently lost.

### Expected behavior

`EnforceDistribution` should preserve the `fetch` value when removing and re-inserting distribution-changing operators.

### Additional context

This is analogous to the existing logic that preserves ordering through `SortPreservingMergeExec` — the `fetch` (limit push-down) should receive the same treatment.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.