fix(datasource): scope file scan work stealing to opened partitions - #173
fix(datasource): scope file scan work stealing to opened partitions#173phillipleblanc wants to merge 4 commits into
Conversation
…ross-partition scan work-stealing DataSourceExec shares one SharedWorkSource queue across all of a scan's partition streams via a per-instance OnceLock, so sibling streams in the same process steal work and each file is opened exactly once. That is correct only when all output partitions execute in the same process. Distributed engines (e.g. Ballista) run each partition in a separate process, where every process rebuilds the full queue over all file_groups and reads the entire input — N-fold scan amplification and, with a Hash-shuffled partial aggregate above the scan, N-fold inflated aggregate results. Add `execution.enable_file_scan_work_stealing` (default true, so single-node behavior is unchanged) and gate `create_sibling_state` in `DataSourceExec::execute` on it. When false, each partition reads only its own file group (`WorkSource::Local`) — the correct behavior under distributed execution. Claude-Session: https://claude.ai/code/session_01G8drJk8Fx5ynRQZ4WkuVju
There was a problem hiding this comment.
Pull request overview
This PR introduces a new execution config option to control cross-partition work-stealing for file scans, addressing correctness issues in distributed execution environments where each scan partition may run in a separate process.
Changes:
- Added
execution.enable_file_scan_work_stealing(defaulttrue) to explicitly control whether sibling scan partitions share a single work queue. - Gated
DataSourceExec::executesibling-state initialization on this option so disabling it forces per-partition local file-group scanning.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
datafusion/datasource/src/source.rs |
Gates creation/usage of per-execution sibling shared state based on the new execution option. |
datafusion/common/src/config.rs |
Adds the execution.enable_file_scan_work_stealing configuration option with detailed rationale and guidance for distributed engines. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Copilot review follow-ups on the `enable_file_scan_work_stealing` gate in `DataSourceExec::execute`: - Document that `create_sibling_state` is the work-stealing hook and is currently only implemented by file-scan sources (`FileScanConfig`); every other `DataSource` inherits the default `None`, so the gate is effectively file-scan scoped and the option name is accurate. Note that a future `DataSource` introducing its own sibling state would also be governed by the flag and should revisit the intended behavior. - Add a regression test that drives `DataSourceExec::execute` with a probe `DataSource` and asserts the gate flips sibling state: shared when `enable_file_scan_work_stealing=true`, absent when `false` (so each partition opens in isolation and reads only its own file group — no cross-partition work-stealing under distributed execution). Claude-Session: https://claude.ai/code/session_01G8drJk8Fx5ynRQZ4WkuVju
|
Addressed the review feedback in 3e97c01:
|
Create SharedWorkSource empty and register each partition's file group only when that partition is opened. This preserves local sibling work stealing while preventing distributed workers that execute a single partition from scanning every file group. Remove the execution.enable_file_scan_work_stealing config gate and add/adjust morsel tests for single-partition registration and late partition registration.
phillipleblanc
changed the title
|
Superseded by the upstream-aligned ballista-side fix. apache/datafusion-ballista#1906 restricts each task's Recommend closing — or repurposing the register-on-open approach as an upstream apache/datafusion proposal if we want the root-cause fix to live in DataFusion. |
Summary
DataSourceExeccreates sibling-shared state once per scan execution so file-scan streams can steal unopened files from one another. The oldSharedWorkSourcewas populated with all file groups up front. That is only safe when the same process opens all output partitions of the scan.Distributed execution engines such as Ballista can run each output partition in a separate process. In that model, every process recreated a shared queue containing the full input and its one local stream could drain every file group. That caused N-fold scan amplification and, with partial aggregation above the scan, inflated aggregate results.
This fixes the ownership model instead of adding a session flag:
SharedWorkSourcenow starts empty, and eachFileStreamregisters only the file group for the partition it opens. A process that opens one partition can only scan that partition's files; a local process that opens all partitions still gets cross-partition work stealing.Changes
SharedWorkSourcefor file scans.FileStreamBuilderopens that partition.preserve_orderandpartitioned_by_file_groupon local work sources, as before.execution.enable_file_scan_work_stealingoption and theDataSourceExecgate from the earlier version of this PR.Test plan
cargo test -p datafusion-datasource --lib