fix: preserve fetch-carrying operators in remove_dist_changing_operators

[v0y4g3r]

Problem

EnforceDistribution::remove_dist_changing_operators strips top-level RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec, assuming they are pure distribution-changing operators that can be regenerated later if necessary.

A fetch on these operators is not a distribution concern — it carries global limit semantics (e.g. produced by LimitPushdown folding a GlobalLimitExec into CoalescePartitionsExec: fetch=N). Stripping the operator silently drops the limit.

This is latent in a single optimizer pass, but breaks when the physical optimizer pipeline runs twice on the same plan (e.g. GreptimeDB plans a nested subplan separately, then optimizes the enclosing plan again — the MergeScan local-execution fallback for information_schema tables): the second pass removes CoalescePartitionsExec: fetch=N and never re-adds it, so

SELECT DISTINCT a FROM t LIMIT 1

can return one row per partition instead of one row in total (GreptimeTeam/greptimedb#8958). Reproducible with vanilla DataFusion 53.1.0 by running the physical optimizer pipeline twice.

Fix

Stop the removal loop at any operator carrying a fetch.

Tests

• keep_fetch_carrying_dist_changing_operators: EnforceDistribution preserves fetch-carrying CoalescePartitionsExec and SortPreservingMergeExec.
• test_global_limit_survives_second_optimizer_pass: running the full physical optimizer pipeline a second time over the optimized plan of SELECT DISTINCT a FROM t LIMIT 1 keeps the global limit boundary (CoalescePartitionsExec: fetch=1 / GlobalLimitExec).

Verified: cargo test -p datafusion --test core_integration 917/917 passed; clippy and fmt clean.

Related: apache#23800 (orthogonal LimitPushdown fix — that one mistakes a per-partition fetch for a global limit; this one drops a global fetch when removing distribution operators).

[Copilot]

Pull request overview

This PR fixes a correctness bug in the physical optimizer where EnforceDistribution::remove_dist_changing_operators could strip distribution-changing operators that also carry fetch (global limit semantics), causing the global limit boundary to be lost when the physical optimizer pipeline is run multiple times on an already-optimized plan.

Changes:

• Stop stripping top-level RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec once an operator with fetch is encountered.
• Add regression coverage ensuring fetch-carrying distribution-changing operators are preserved by EnforceDistribution.
• Add an integration regression test ensuring SELECT DISTINCT a ... LIMIT 1 retains a global limit boundary after a second optimizer pass.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
datafusion/physical-optimizer/src/enforce_distribution.rs	Preserve fetch-carrying distribution-changing operators during the “strip top operators” phase to avoid dropping global limit semantics.
datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs	Add regression test that rerunning physical optimization keeps the global limit boundary for a limited DISTINCT query.
datafusion/core/tests/physical_optimizer/enforce_distribution.rs	Add targeted unit test that EnforceDistribution keeps fetch-carrying CoalescePartitionsExec and SortPreservingMergeExec.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.