[WIP]feat(aqe): port Spark's OptimizeSkewedJoin

[wirybeaver]

Summary

Ports Spark's OptimizeSkewedJoin AQE rule into Ballista. When a shuffle partition is disproportionately large (skewed), the rule splits it into smaller sub-ranges and replicates the matching partition on the other side of the join, so the work that was a single straggler task becomes N parallel tasks.

Replaces the file-list sharding approach (#1718), which had to bail on any HashPartitioned / SinglePartition consumer — effectively all joins and FinalPartitioned aggregates.

Design

Spark	Ballista
OptimizeSkewedJoin	OptimizeSkewedJoinRule in state/aqe/optimizer_rule/
MapOutputStatistics	ExchangeExec::shuffle_partitions() + partition_stats.num_bytes()
PartialReducerPartitionSpec(reducerIdx, [startMap, endMap))	SkewJoinShard { upstream_idx, start_map_idx, end_map_idx }
Cartesian pair-up across legs	Pre-computed Vec<SkewJoinShard> of length N on both legs

Rule is opt-in: SET ballista.planner.skew_join.enabled = true. Defaults match Spark (factor=5.0, threshold=256 MiB, advisory=64 MiB).

Current limitations

• v1 only handles single-join stages. Bails when a stage subtree has more than one binary join. Per-join iteration is future work.
• Mutually exclusive with DataFusion's dynamic-filter pushdown. The skew rewrite breaks the hash-colocation invariant that dynamic filters assume. A guard enforces SET datafusion.optimizer.enable_dynamic_filter_pushdown = false.
• Split granularity is per upstream task, not per row group. Single-file tables produce 1 upstream task per partition — nothing to subdivide. Tables must span multiple files. (Spark doesn't have this limitation because it splits large files into multiple scan tasks automatically.)

---

End-to-end verification

Setup

cd /home/user/optimize-skewed-join
cargo build --release -p ballista-scheduler -p ballista-executor -p ballista-cli

# Scheduler with AQE debug logging
RUST_LOG=info,ballista_scheduler::state::aqe=debug \
  ./target/release/ballista-scheduler --bind-port 50050 > /tmp/scheduler.log 2>&1 &
sleep 3

# Executor
RUST_LOG=info \
  ./target/release/ballista-executor \
    --scheduler-host localhost --scheduler-port 50050 \
    --bind-port 50051 --concurrent-tasks 4 > /tmp/executor.log 2>&1 &

Generate synthetic skewed data — 20 files, ~90% of rows have key=0:

mkdir -p /tmp/skewed_dir
for i in $(seq 0 19); do
  duckdb -c "
    COPY (
      SELECT (random() < 0.9)::int AS key, range AS val
      FROM range(50000)
    ) TO '/tmp/skewed_dir/part_${i}.parquet' (FORMAT parquet);
  "
done

duckdb -c "
  COPY (SELECT range AS key, range AS val FROM range(100))
  TO '/tmp/dim.parquet' (FORMAT parquet);
"

Note: ballista-cli must be invoked with --host localhost --port 50050. Without these flags it runs queries in-process and the scheduler is never contacted.

Baseline (rule OFF)

cat > /tmp/run-baseline.sql << 'EOSQL'
SET ballista.planner.adaptive.enabled = true;
SET ballista.planner.skew_join.enabled = false;
SET datafusion.optimizer.enable_dynamic_filter_pushdown = false;

CREATE EXTERNAL TABLE fact STORED AS PARQUET LOCATION '/tmp/skewed_dir/';
CREATE EXTERNAL TABLE dim  STORED AS PARQUET LOCATION '/tmp/dim.parquet';

SELECT dim.val, count(*) FROM fact JOIN dim ON fact.key = dim.key GROUP BY dim.val;
EOSQL

./target/release/ballista-cli --host localhost --port 50050 \
  -f /tmp/run-baseline.sql > /tmp/baseline.out 2>&1

Rewritten (rule ON)

cat > /tmp/run-rewritten.sql << 'EOSQL'
SET ballista.planner.adaptive.enabled = true;
SET ballista.planner.skew_join.enabled = true;
SET ballista.planner.skew_join.skewed_partition_factor = 1.5;
SET ballista.planner.skew_join.skewed_partition_threshold_bytes = 1024;
SET ballista.planner.skew_join.advisory_partition_bytes = 1024;
SET datafusion.optimizer.enable_dynamic_filter_pushdown = false;

CREATE EXTERNAL TABLE fact STORED AS PARQUET LOCATION '/tmp/skewed_dir/';
CREATE EXTERNAL TABLE dim  STORED AS PARQUET LOCATION '/tmp/dim.parquet';

SELECT dim.val, count(*) FROM fact JOIN dim ON fact.key = dim.key GROUP BY dim.val;
EOSQL

./target/release/ballista-cli --host localhost --port 50050 \
  -f /tmp/run-rewritten.sql > /tmp/rewritten.out 2>&1

Verify

# Results match
diff <(grep -E '^\|' /tmp/baseline.out | sort) \
     <(grep -E '^\|' /tmp/rewritten.out | sort)

# Rule attached
grep -a 'attaching SkewJoinPlan' /tmp/scheduler.log

Results

attaching SkewJoinPlan: M=16 K'=34
  left_skewed=[…, true, …, true, …]   (partitions 3 and 5)
  left_sizes=[…, 10817136, …, 1207824, …]
  right_skewed=all false               (dim table — 100 rows)

• Correctness: rows identical between baseline and rewritten.
• Skew detection: 2 of 16 partitions flagged (10.8 MB and 1.2 MB vs near-zero for the rest).
• Split: 34 downstream join tasks instead of 16. Non-skewed side replicated for each sub-range.
• Mutual-exclusion guard: verified separately — enable_dynamic_filter_pushdown=true triggers bail as expected.

Supersedes #1718.

🤖 Generated with Claude Code

[martin-g]

@wirybeaver Please resolve the conflicts!
Update: Sorry, I didn't notice that it is still a draft!