apache / apache/datafusion-ballista

Introduce TaskGroups to take advantage of shared memory between task slots

Open
#332 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Rust
Stars
2.1k
Forks
320
Avg merge
1d 22h
Merged PRs (30d)
66

Description

**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
(This section helps Arrow developers understand the context and *why* for this feature, in addition to the *what*)

Currently, the Ballista scheduler treats each task slot as effectively its own machine but in practice an executor may have one task slot per CPU core. So in a cluster with 12 executor nodes each with 12 cores, the scheduler would treat them like 144 independent machines. This adds a lot of unnecessary inefficiency to some common queries.

To take a simple example, for a query like `SELECT * FOM table WHERE foo = 'bar' LIMIT 1000` on 144 partitions, we would break this into a two stage query

```
Stage 2:
GlobalLimitExec: fetch=1000
ShuffleReaderExec

Stage 1:
ShuffleWriterExec
LocalLimitExec
FilterExec
ParquetExec
```

Each partition would read 1000 rows and then we would schedule the second stage which would stop after reading one of the 144 output partitions from the previous stage. So we've effectively read 144k rows from when we only need 1k.

While it's not possible to eliminate this issue entirely in a distributed query we can do better than the current approach I think.

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

If the scheduler is able to schedule multiple partitions on the same node (eg executor slots on a single executor) then it can rewrite the plan to coalesce those tasks into a `TaskGroup` which can be executed more efficiently with existing constructs in DataFusion. For example. for our example query above, if all 12 task slots are available on a given executor, then we can take those twelve tasks and scheduler them as a single execution plan:

```
GlobalLimitExec: fetch=1000
CoalescePartitionsExec
LocalLimitExec: fetch = 1000
FilterExec
ParquetExec
```

This can be wrapped in a `TaskGroup` data structure that (simplified) looks something like:

```rust
struct TaskGroup {
partitions: Vec,
plan: Arc
}
```

Then the stage 2 plan needs to be adjusted since it will read a smaller number of input partitions.

Alternatively we can generalize the existing data structures to handle this case.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

We could leave things as is

**Additional context**
Add any other context or screenshots about the feature request here.

Contributor guide

Open the contributing guide

Research direction

Start by tracing the scheduler's task-slot handling and the stage plans described in the issue, then review how DataFusion's CoalescePartitionsExec and execution plans could represent grouped partitions. Done means task slots on one executor can be scheduled as a TaskGroup and the downstream stage reads the reduced number of input partitions, with the limit query behavior checked end to end.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.