Skip to content

Comet Native Executors: use Ballista's adaptive planning for joins inside Comet fragments #7391

Open

Description

@andygrove

What is the problem the feature request solves?

On a Comet Native Executor cluster (#7289), a query gets none of Ballista's adaptive planning. Where Comet on Spark switches a join to a broadcast at runtime, Comet on Ballista runs Spark's static plan.

TPC-H SF1000 on the same cluster shape (32 executors of 8 vCPU and 64 GiB, Amazon EKS), mean of 2 iterations:

Pipelined Blocking
Comet on Ballista (#7289) 462.2 s 593.6 s
Ballista (main 0ef7399e) 281.1 s 483.9 s
Comet 1.1.0-rc2 on Spark (published on Ballista's benchmarking page, another run) 455.1 s

That mixes two gaps:

  • Comet on Ballista is slower than Comet on Spark on 8 queries, by 53.8 s in all. q18 (+32.0 s) and q17 (+15.7 s) are most of it. Comet on Spark gets there through Spark's AQE.
  • Comet on Spark is itself slower than Ballista on 18 queries, by 181.6 s in all (q9, q8, q5 and q21 lead). That is Spark's plan against DataFusion's, and no adaptive rule closes it.

q17 shows the first one. The scheduler log has the stage graph that ran: lineitem (2364 scan partitions) is hash-shuffled on l_partkey twice, part is hash-shuffled on p_partkey, and one 512-partition fragment joins the three. No input is a broadcast. It takes 61 s. Ballista's plan for the same query filters part to 199,303 rows, broadcasts it to the lineitem scan, and so shuffles 5,980,497 rows of lineitem for that join, not 6 billion. It takes 22 s (blocking) or 14 s (pipelined).

Why Comet on Ballista gets none of it (from native/native-executors/src/assembly.rs and Ballista's comet-native branch):

  1. The joins are inside CometFragmentExec blocks, serialized Comet Operator trees. Ballista's adaptive rules look for DataFusion join and exchange nodes, so they see a leaf.
  2. Spark's exchanges are already placed. A hash exchange becomes a RepartitionExec, which PinSubmittedShufflesRule turns into an ExchangeExec so DataFusion's optimizers cannot remove it. That keeps Spark's plan, including a shuffled join where a broadcast would do.
  3. Coalescing is switched off for a shuffle that a Comet fragment reads.
  4. A broadcast is decided before the query runs: StageInput.broadcast comes from Spark's static plan.
  5. Pipelined mode, the faster one, has no adaptive planning at all, since it places all tasks up front.
  6. The rule list is fixed (AdaptivePlanner::default_optimizers is pub(super), and try_new_from_physical takes no rules), so comet-scheduler cannot add one without a Ballista change.

Describe the potential solution

In order of what I would try:

  • A. A rule for Comet fragments in Ballista's adaptive planner (blocking mode). When the input exchange of a fragment that holds a join has finished below the broadcast threshold, set that input's broadcast flag and drop the other side's exchange, so the large side is read in place. It needs a way to add a rule (a hook in Ballista, or a trait that a custom node implements to say "I am a join: type, keys, which input is which") and the join's description from the driver in Stage. The rule must only broadcast where it is correct: the right side for inner, left semi, left anti and left outer, never for right or full outer or a null-aware anti join with a nullable key. It does not help pipelined mode.
  • B. Lift the joins out of the blocks as DataFusion HashJoinExec/SortMergeJoinExec with plain column keys and opaque fragments below, so all of Ballista's rules (and skew handling) apply. Larger, and a join with a non-equi filter keeps a Comet expression and stays opaque.
  • C. Spark's AQE on the driver, stage by stage. All of Spark's rules, with no new rule code, but later jobs must read earlier jobs' output from the executors, every stage becomes a barrier, and pipelined mode's overlap is gone.
  • D. Better estimates, or cheap probes, for pipelined mode. Pipelined mode cannot adapt. Either give Spark's static plan statistics it lacks (a path-based Parquet view has none, so a filter does not shrink Spark's estimate), or run small jobs for the filtered build sides that would otherwise be shuffled (q17's part stage took 0.7 s of compute) and plan a broadcast when they come back small.

Before any code:

  1. Run Comet on Spark on q17 with spark.sql.adaptive.enabled=false. If the diagnosis is right it slows from about 25 s to about 60 s.
  2. Read the Comet event logs for q17 and q18 to see how many rows each stage shuffles.
  3. Print Spark's size estimate for the filtered part (EXPLAIN COST).

Additional context

Not verified: that Comet on Spark's AQE takes the same step on q17 (step 1 tests it); that Ballista's planner re-runs its rules when the stage under a CometFragmentExec finishes (only the join and exchange rules were read); and how much of Ballista's pipelined lead on q5, q8 and q9 is join order and how much is broadcast choice.

Ballista follow-ons if A is the route: a way to add rules to the adaptive planner, in apache/datafusion-ballista.

Part of #7288.

Willingness to contribute

Yes. I can contribute this feature independently.

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions