Skip to content

fix: preserve memory scan fetch during repartitioning - #25068

Closed
discord9 wants to merge 1 commit into
apache:mainfrom
discord9:fix/preserve-memory-source-fetch
Closed

fix: preserve memory scan fetch during repartitioning#25068
discord9 wants to merge 1 commit into
apache:mainfrom
discord9:fix/preserve-memory-source-fetch

Conversation

@discord9

@discord9 discord9 commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

No linked issue. Independently reproduced physical-optimizer correctness fix.

Rationale for this change

Repartitioning a memory source with an existing fetch rebuilds the source without that fetch. A one-partition scan containing batches [0,1] and [2,3] with fetch 1 returns [0] before optimization, but all four rows after source repartitioning. This also occurs through EnsureRequirements and one/two complete default physical-optimizer passes with a parent filter and no outer limit.

The fetch is enforced separately for each original partition. Copying it onto newly formed partitions would still change the result, rather than preserve the original cap.

This is demonstrated for directly constructed physical plans and their optimization; an ordinary SQL-planning reproduction is not established.

What changes are included in this PR?

Have MemorySourceConfig decline source-level repartitioning whenever a fetch is present, before either sorted or unsorted repartitioning. The optimizer can still insert an ordinary RepartitionExec over the existing capped source, distributing rows already admitted by the original caps.

This uses the existing optional repartition contract, with no changes to optimizer configuration, statistics, or shared fetch accounting. It is independent of #23800 and #25065.

What is the testing strategy for this PR?

  • Uncapped source repartition remains enabled and emits all rows.
  • Eligible two-batch sources with fetch 0 decline repartition and remain empty.
  • Two original capped partitions preserve their prefix values [0,10].
  • An eligible projected/sorted source preserves its schema, ordering, and capped value [10].
  • EnsureRequirements and one/two full optimizer passes retain [0] and use ordinary round-robin repartition over the original one-partition fetched source.

Removing only the guard makes the independently eligible projection/ordering regression fail. The guard was restored byte-exactly; the original direct/API and optimizer reproductions also demonstrated the pre-fix one-row-to-four-row change.

Verified locally:

  • cargo check -p datafusion-datasource
  • cargo test -p datafusion-datasource memory:: --lib — 15 passed
  • cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution:: — 88 passed
  • cargo fmt --all --check
  • cargo clippy --workspace --all-targets --all-features -- -D warnings

The full workspace runtime test suite was not run.

Are there any user-facing changes?

Preserves the results of already-capped memory scans during physical optimization. No public API or configuration changes. Source-level repartition is declined for capped sources; ordinary downstream repartition remains available.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
@github-actions github-actions Bot added core Core DataFusion crate datasource Changes to the datasource crate labels Sep 8, 2026
@discord9 discord9 closed this Sep 8, 2026
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 83.00654% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.74%. Comparing base (b39cb4a) to head (5997d8d).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/datasource/src/memory.rs 83.00% 3 Missing and 23 partials ⚠️
Additional details and impacted files
@@           Coverage Diff            @@
##             main   #25068    +/-   ##
========================================
  Coverage   81.74%   81.74%            
========================================
  Files        1128     1128            
  Lines      416644   416797   +153     
  Branches   416644   416797   +153     
========================================
+ Hits       340585   340712   +127     
- Misses      55999    56003     +4     
- Partials    20060    20082    +22     

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate datasource Changes to the datasource crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants