fix: preserve fetch across distribution reoptimization

[xudong963]

Which issue does this PR close?

• Closes EnsureRequirements can silently drop fetch during distribution reoptimization #24807.

Rationale for this change

Reoptimizing an already optimized physical plan could silently remove a pushed-down LIMIT stored on SortPreservingMergeExec or CoalescePartitionsExec. This could make a query return more rows than requested.

What changes are included in this PR?

• Consume a removed fetch when a replacement merge operator is inserted.
• Remember the outermost removed fetch-capable distribution operator and rebuild it around the optimized child when no replacement consumes its limit.
• Preserve the minimum effective fetch across nested distribution operators.
• Carry a fetched ordered merge's limit to a replacement sort when order-preserving variants are removed.

What is the testing strategy for this PR?

Added targeted physical optimizer regression tests covering:

• reoptimizing a fetched SortPreservingMergeExec;
• reoptimizing a fetched CoalescePartitionsExec;
• moving a fetched ordered merge's limit to a replacement sort;
• updating an existing fetched single-partition merge snapshot that previously
  encoded the incorrect removal of its limit.

Validated with:

# Fails on e4cf35cbc (current main before this patch): both regression tests fail
cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::preserve_fetch_when_reoptimizing

# Passes with this patch
cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::preserve_fetch_when_reoptimizing
cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::move_fetch_to_replacement_sort
cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::test_replace_order_preserving_variants_with_fetch
cargo fmt --all -- --check
cargo clippy --all-targets --all-features -- -D warnings
RUST_BACKTRACE=1 cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption

Are there any user-facing changes?

Queries preserve their requested global limit when physical distribution requirements are optimized more than once. There are no changes to SQL behavior other than fixing the incorrect result, and no changes to documented public APIs.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.72%. Comparing base (84ccbb1) to head (bd6dfd1).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24809      +/-   ##
==========================================
- Coverage   81.72%   81.72%   -0.01%     
==========================================
  Files        1127     1127              
  Lines      416395   416508     +113     
  Branches   416395   416508     +113     
==========================================
+ Hits       340306   340392      +86     
- Misses      56098    56114      +16     
- Partials    19991    20002      +11     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[jayzhan211]

Thanks @xudong963 , here is a suggestion:

The ordering_satisfied branch of replace_order_preserving_variants_with_fetch detaches the fetch from the merge and re-attaches it on the sort added above the whole child context. That is only equivalent when everything between the merge and the sort is row-preserving, and the operators that keep the SPM alive until the parent's visit (maintains_input_order == true) are not all row-preserving.

Repro on this branch, EnsureRequirements::optimize with enable_round_robin_repartition = false:

input:  SortRequiredExec -> FilterExec: c@2 > 0 -> SortPreservingMergeExec: [c@2 ASC], fetch=5 -> DataSourceExec(2 sorted partitions)
output: SortRequiredExec -> SortPreservingMergeExec: [c@2 ASC], fetch=5 -> FilterExec: c@2 > 0 -> DataSourceExec

The input filters the 5 smallest rows; the output takes the 5 smallest filtered rows. Different result set.

Same root cause, second symptom: the fetch is only materialised if add_sort_above_with_check actually adds a sort. With the test file's filter_exec (predicate c = 0, which makes c constant so the ordering is trivially satisfied) the output is SortRequiredExec -> FilterExec -> CoalescePartitionsExec -> DataSourceExec with no fetch anywhere, which is the bug this PR is meant to close.

A fetched SPM is a TopK merge, so replace it in place with a TopK sort over the coalesce and never hand the fetch back to the caller. That also lets the Option<usize> return value and the min_fetch(preserved_fetch, output_fetch) call in ensure_distribution go away:

fn replace_order_preserving_variants_impl(
    mut context: DistributionContext,
    ordering_satisfied: bool,
) -> Result<DistributionContext> {
    context.children = context
        .children
        .into_iter()
        .map(|child| {
            if child.data {
                replace_order_preserving_variants_impl(child, ordering_satisfied)
            } else {
                Ok(child)
            }
        })
        .collect::<Result<Vec<_>>>()?;

    if let Some(spm) = context.plan.downcast_ref::<SortPreservingMergeExec>() {
        let child_plan = Arc::clone(&context.children[0].plan);
        context.plan = match spm.fetch() {
            // A fetched merge is a TopK. Keep the limit at this position by
            // replacing it with a TopK sort over the coalesced input; the
            // sort also satisfies the ordering the merge provided.
            Some(fetch) if ordering_satisfied => Arc::new(
                SortExec::new(spm.expr().clone(), Arc::new(CoalescePartitionsExec::new(child_plan)))
                    .with_fetch(Some(fetch)),
            ),
            fetch => Arc::new(CoalescePartitionsExec::new(child_plan).with_fetch(fetch)),
        };
        return Ok(context);
    } else if let Some(repartition) = context.plan.downcast_ref::<RepartitionExec>()
        && repartition.preserve_order()
    {
        // unchanged
    }

    context.update_plan_from_children()
}

and in ensure_distribution:

-                        let (replaced_context, preserved_fetch) =
-                            replace_order_preserving_variants_with_fetch(context, ordering_satisfied)?;
-                        context = replaced_context;
+                        context = replace_order_preserving_variants_impl(context, ordering_satisfied)?;
                         if ordering_satisfied {
-                            let output_fetch = ...;
-                            context = add_sort_above_with_check(context, sort_req, min_fetch(preserved_fetch, output_fetch))?;
+                            context = add_sort_above_with_check(context, sort_req, output_fetch)?;
                         }

move_fetch_to_replacement_sort still passes with this shape (the TopK sort ends up directly above the coalesce). Please add the two shapes above as regression tests: fetch must stay below the FilterExec, and must survive when no extra sort is needed.

[kosiew]

@xudong963,

Thanks for working on this. The added coverage around preserving fetch across reoptimization is helpful. I found one correctness issue with nested fetched distribution operators that I think needs to be addressed before merging. I also left one non-blocking suggestion to strengthen the TopK regression test.

[kosiew]

I think there is still a correctness issue when multiple removed distribution operators have fetches. removed_fetch collapses them to the minimum value, while fetch_plan remembers only the outermost fetched operator. That loses the semantic position of the inner fetch.

For example, consider CoalescePartitionsExec(fetch=10) -> SortPreservingMergeExec([c], fetch=5) -> two sorted partitions. The original plan gets the global TopK 5 from the ordered merge. After both operators are removed, we retain fetch=5 but can restore it as CoalescePartitionsExec(fetch=5) directly over the partitions. That can return the first five rows in coalesce/input order rather than the global TopK 5.

Could we preserve each fetched operator at its original semantic boundary, or otherwise replace it with something that is provably equivalent? I think it would also be useful to add an execution regression for this nested Coalesce/SPM case, using partition values where concatenation order differs from global sort order.

[xudong963]

Thanks for the concrete example. Fixed in 8098706: distribution cleanup now stops at fetched operators, keeping each limit at its original position instead of collapsing fetch values and restoring only the outermost operator. Replacing a fetched SPM also retains its TopK at that position with the original merge ordering, even when the parent requires a different ordering.

The nested execution regression uses even/odd partitions and outer fetch values of 0, 3, and 10 around an SPM fetching 5. It checks exact results before optimization and after two passes, with sort parallelization both enabled and disabled.

I also protected fetched Coalesce nodes from removal during sort parallelization and replacement with ordered merges. Added regressions for reversed ancestor ordering and a filter above a fetched coalesce. All three new execution regressions fail on c90e58b and pass with this fix.

[kosiew]

Could we strengthen this test by executing a small two-partition input and asserting the resulting TopK values as well? The display assertion confirms that we constructed a SortExec with fetch=5, but an execution assertion would also protect the actual ordering, null handling, tie behavior, and fetch placement if the implementation changes later. This is non-blocking.

[xudong963]

Done in 8098706. move_fetch_to_replacement_sort now executes two sorted partitions and compares the actual TopK values before and after replacement, while still checking the replacement sort's fetch=5.

The cases cover ascending order with NULLs first, descending order with NULLs last, and duplicate values. This checks the resulting values and ordering as well as the plan shape.

[xudong963]

Thanks @jayzhan211, addressed in c90e58b and 8098706. A fetched SPM is now replaced in place with a TopK sort over the coalesce, using the merge's original ordering. The helper no longer hands a detached fetch back to the caller. The coalesce also has its own DistributionContext to keep the context tree consistent with the plan.

Added executable regressions for both c > 0 and c = 0, checking actual results before optimization and after two optimizer passes, plus context integrity. Both fail on the original PR head and pass with the fix, covering fetch placement below the filter and the case where no extra sort is needed.

[kosiew]

@xudong963,

Thanks for the follow-up. I went through the changes again and the issues from the previous review look addressed.

The fetched CoalescePartitionsExec and SortPreservingMergeExec nodes now remain at their semantic boundaries, so nested fetches are not collapsed or moved to a different point in the plan. The fetched SPM replacement also looks correct: materializing it as a TopK SortExec over the coalesced input preserves both the ordering and fetch before any ancestor operators.

The additional execution coverage is helpful as well. In particular, the nested even/odd partition case exercises the original correctness issue through repeated optimizer passes, and the TopK tests cover ascending and descending ordering, NULLs, and duplicate values. The filter regression also verifies that fetch does not move across the filter during reoptimization.

I also checked the sort parallelization and order-preserving rewrite paths. Fetched coalesces are now protected from being replaced or removed, which addresses the remaining dropped-fetch cases discussed in the review.

I did not find any new correctness issues in the follow-up changes. Thanks for addressing the feedback.

[xudong963]

Also CC @zhuqi-lucas in case you have time to have a look

[zhuqi-lucas]

LGTM Thanks @xudong963 !