apache / apache/datafusion-ballista
Pinot-style colocated-join optimizer for hash-bucketed tables
- Dominant language
- Rust
- Stars
- 2.1k
- Forks
- 320
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 66
Description
## Problem
Ballista shuffles data on every join whose keys don't already match the input partitioning, because the planner has no way to learn that a table is hash-bucketed by some column. For workloads on pre-bucketed data (e.g. Iceberg / Parquet datasets where producers already hash-partition by the join key), the network shuffle is pure overhead: the data is *already* co-located, the planner just doesn't know it.
Concretely:
- `BallistaQueryPlanner` ships the logical plan to the scheduler.
- `AdaptivePlanner` runs DataFusion's physical planner + a custom rule chain ending in `DistributedExchangeRule`, which inserts an `ExchangeExec` above every `RepartitionExec(Hash)`.
- `DefaultDistributedPlanner` cuts a stage at every shuffle.
- The only join-aware rule today is `JoinSelection`, which only swaps build sides based on row/byte stats — it has no notion of partition co-location.
- The catalog has no concept of hash bucketing; only `ListingTable` directory partitions exist.
## Proposed solution
Bring three Pinot V2 Physical Optimizer ideas to Ballista (Pinot's `RelToPRelConverter` → `WorkerExchangeAssignmentRule`):
1. **Colocated joins** — when both join inputs declare matching hash bucketing on the join keys, skip the shuffle entirely and let the join's required hash distribution be satisfied directly by the inputs.
2. **Sub-partitioning** for divisor bucket counts (e.g. 16 vs 8) — locally coalesce the larger side instead of shuffling, since `(hash(k) % 16) % 8 == hash(k) % 8`.
3. **Small-side broadcast** in the AQE path — when one side fits under a configured byte threshold, promote a `Partitioned` `HashJoinExec` to `CollectLeft`. (This complements PR #1647, which added the same lowering to the non-AQE `DefaultDistributedPlanner`; the AQE planner has no equivalent today, per the TODO at `ballista/scheduler/src/state/aqe/mod.rs`.)
## Architecture
A small metadata trait in `ballista-core` lets any `TableProvider` declare on-disk hash bucketing without a DataFusion fork. Two new `PhysicalOptimizerRule`s slot into the existing `AdaptivePlanner` rule chain just before `DistributedExchangeRule`:
```
DataFusion physical planner
→ DataFusion optimizer rules (incl. EnforceDistribution, JoinSelection)
→ ColocatedJoinRule ← strip redundant RepartitionExec
(also handles divisor sub-partitioning)
→ BroadcastSmallSideRule ← promote Partitioned join to CollectLeft
→ DistributedExchangeRule ← existing: maps remaining repartitions to ExchangeExec
→ DefaultDistributedPlanner ← existing: cuts stages at ExchangeExec
```
## Goals
- Let users declare per-table hash partitioning so the optimizer can reason about it.
- Skip the shuffle when both join inputs are already hash-partitioned on the join keys with matching bucket counts.
- When bucket counts differ but have a divisor relationship, use a sub-partitioning exchange that keeps all data local.
- When one side is below a configurable size threshold, broadcast it instead of repartitioning the other side — for the AQE code path that PR #1647 explicitly left as a TODO.
## Non-goals
- A full PRel-style multi-pass optimizer rewrite (reuse DataFusion's `Partitioning` + `EquivalenceProperties` instead of inventing a `BallistaDataDistribution`).
- Custom DDL grammar — piggyback on DataFusion's existing `OPTIONS (...)` clause and a `PartitionedTableProvider` wrapper.
- Iceberg / Delta integration (design the trait so they can plug in later).
- Rewriting `JoinSelection`'s cost model.
## Implementation
PR #1676.
## Follow-up
#1679 — extend `ColocatedJoinRule` and `BroadcastSmallSideRule` to `SortMergeJoinExec`, since PR #1651 made sort-merge the default join.
Contributor guide
Research direction
Start by reading the AdaptivePlanner rule chain and ballista/scheduler/src/state/aqe/mod.rs, then review PR #1676 to understand the proposed colocated-join and broadcast behavior. Done means the declared bucketing metadata is usable by the optimizer, redundant repartitions and eligible AQE broadcasts are handled, and the documented goals are covered without taking on the SortMergeJoin follow-up in #1679.
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
- Clearly specified
- Newbie friendliness
- 25/100