Pinot-style colocated-join optimizer for hash-bucketed tables

[wirybeaver]

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 feat(scheduler): broadcast-style hash join for small-side joins #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 PhysicalOptimizerRules 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 feat(scheduler): broadcast-style hash join for small-side joins #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.