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

❌ Patch coverage is 93.54839% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.69%. Comparing base (e4cf35c) to head (c90e58b).
⚠️ Report is 90 commits behind head on main.

Files with missing lines	Patch %	Lines
...er/src/ensure_requirements/enforce_distribution.rs	93.54%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24809      +/-   ##
==========================================
+ Coverage   81.52%   81.69%   +0.17%     
==========================================
  Files        1123     1127       +4     
  Lines      406148   415590    +9442     
  Branches   406148   415590    +9442     
==========================================
+ Hits       331124   339533    +8409     
- Misses      55659    56116     +457     
- Partials    19365    19941     +576     

☔ 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.

[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.