fix: Enforce data distribution requirements and fix plans that break them

[barbarj]

Closes #563

The bug in 563 is actually a failure to enforce a plan shape invariant for multi-task stages: for every stage, executing its plan once per task over the per-task input assignment and unioning the outputs must be equivalent to executing it once over all the data.

This PR adds a correctness check after plan finalization (validate_stages.rs) and applies fixes for the problematic shapes (normalize_collect_joins.rs).

This runs validate_distributed_stages() as the final step before returning during physical plan construction. We classify the dataflow across each stage, bottom-up, as partitioned or replicated. Partitioned data flow is safe across multiple tasks. Replicated is not. In cases where the input to a task is replicated data, we (based on our knowledge of node semantics) convert that to partitioned if the join will not emit build-side rows, making the replication (i.e. the broadcast build-side) safe as input and the output meaningfully partitioned. See validate_stages.rs for a more in-depth explanation and justification.

There are 5 invalid plan shapes to fix. The following strategies are used:

• HashJoin CollectLeft (join types: Left/LeftSemi/LeftAnti/LeftMark/Full) → change
  CollectLeft->Partitioned (keys exist; no replication needed)
• HashJoinCollectLeft, null_aware (join type: LeftAnti) -> cap at 1 task
• NLJ (join types: Left/LeftSemi/LeftAnti/LeftMark) -> call swap_inputs(), then broadcast as usual. (makes the join shape safe for broadcast).
• NLJ (join type: Full) -> cap at 1 task
• CrossJoin (always safe with broadcast) -> cap at 1 task if broadcasts are disabled.

The CollectLeft->Partitioned change causes the rewriting of a bunch of plan shapes in the plan tests. The problematic stages run over small data, and thus collapse to a single task during
prepare_network_boundaries, avoiding the bug. Due to the shape normalization pass necessarily running before task counts are decided, we can't see that the task consolidation will happen and must modify the plan.

In-repo benchmarks show equivalent or better performance for all queries, save one: TPCH q22. It drops by about 40% (~17ms -> ~24ms on my machine, 8 workers, 2 threads). As far as I can tell, it's losing a dynamic filter due to being cut into more stages.

[barbarj]

My apologies for not surfacing the validation design before cutting this PR. If you think that should be done differently, I'm happy to limit this PR to just the shape fixes and do the validation separately.

[gabotechs]

Hi @barbarj! thanks so much for this quality PR 🙏.

I'd like to spend some time understanding the problem so that I can better review this and give good suggestions. In the meantime, I think we can start shipping chunks of this PR, for example, I'd like to propose creating a PR with the new tests/multi_task_collect_join_repros.rs integration tests (#[ignored] for now as I imagine they would not pass), would you be fine with that?

[gabotechs]

Flushing thoughts as I read:

This sounds like something that could be useful if we wanted to dynamically swap JOIN orders at runtime during AQE (Adaptive Query Execution):

#522

[gabotechs]

Just ran the remote benchmarks against main on a 12 machine cluster:

tpch_sf1
   TOTAL: prev=5216 ms, new=4649 ms, diff=1.12 faster ✔
tpch_sf10
   TOTAL: prev=12052 ms, new=11218 ms, diff=1.07 faster ✔
tpch_sf100
   TOTAL: prev=59778 ms, new=59279 ms, diff=1.01 faster ✔
tpcds_sf1
   TOTAL: prev=35592 ms, new=35534 ms, diff=1.00 faster ✔
clickbench_0-100
   TOTAL: prev=28801 ms, new=29263 ms, diff=1.02 slower ✖

So performance wise looks good 👍

[gabotechs]

Flushing some more comments as I keep reviewing.

[gabotechs]

Left some comments for the validate_stages.rs file.

The normalize_collect_joins.rs and test from this PR are mostly ready to go, with just some minor cosmetic comments. For shipping those eariler, I'd suggest to split this PR in two:

1. normalize_collect_joins.rs and the tests.
2. validate_stages.rs

I'd like for other people to put some more eyes on 2), but 1) is something we can start shipping sooner. WDYT?

[gabotechs]

I don't think this assumption is right. There are valid cases for NetworkCoalesceExec to be placed in the middle of the plan.

Take for example this:

cargo run \
--features integration \
--example custom_distributed_partial_reduction_tree \
'SELECT "RainToday", count(*) FROM weather GROUP BY "RainToday"' \
--show-distributed-plan

The runnable example no longer works in this PR, even though it's a valid distributed plan.

[barbarj]

I see your point. NetworkCoalesceExec maintains it's input partitioning, so I think under this model it'd be safe to return DataFlow::Partitioned { claims_are_global: true } here.

[jayshrivastava]

I just started reviewing but here's a random comment

We often have plans where there's a RepartitionExec Hash(expr, M * N) at the head of a stage and the stage above has M tasks of N partitions each. It begs the question - is Partitioning=Hash(expr, M x N) the same as Partitioning=Hash(expr, N)? The inter-task partitioning is Partitioning=Hash(expr, N), however the intra-task partitioning is Partitioning=Hash(expr, N*M). All the tasks pass so it's safe but I could see this becoming a problem / annoyance later.

[jayshrivastava]

The runnable example no longer works in this PR, even though it's a valid distributed plan.

Should we have CI run all our examples? Or port this example into a test. This seems important.

[gabotechs]

Why would this be specific to DistributedLeafExec? Any user can create their own leaf node and claim a Partitioned::Hash output without using DistributedLeafExec.

[barbarj]

As I understand it, it's this property:

// A DistributedLeafExec resolves to a different slice of the underlying source in every task.

which means a Partitioning::Hash implies claims_are_global. Though I suppose any node claiming Partitioning::Hash implies its output is distributed accordingly.

[barbarj]

Re: your below comment:

I feel like the code that handles children.is_empty() and node.is::() should be the same, and should not really be taking into account if there are desired task count handlers registered.

Assuming that I'm right about the task count propagation I mentioned, then this block actually becomes something like:

if node.is::<DistributedLeafExec>() || node.children().is_empty() {
    let claims_are_global = matches!(node.output_partitioning(), Partitioning::Hash(..));
    return Ok(DataFlow::Partitioned { claims_are_global })
}

[gabotechs]

🤔 I'm not sure why this would be needed here. At the moment this function is called, all the task count have already been propagated and they are set in stone, so it matters little whether there are desired task count handlers registered for that node.

[gabotechs]

I feel like the code that handles children.is_empty() and node.is::<DistributedLeafExec>() should be the same, and should not really be taking into account if there are desired task count handlers registered.

[barbarj]

🤔 I'm not sure why this would be needed here. At the moment this function is called, all the task count have already been propagated and they are set in stone, so it matters little whether there are desired task count handlers registered for that node.

This mirrors the how inject_network_boundaries decides task counts:

// from _inject_network_boundaries()

if plan.children().is_empty() {
        // This is a leaf node, maybe a DataSourceExec, or maybe something else custom from the
        // user. We need to estimate how many tasks are needed for this leaf node, and we'll take
        // this decision into account when deciding how many tasks will be actually used.
        let ev = DesiredTaskCountEvent {
            plan: &plan,
            session_config: nb_ctx.cfg,
        };
        return if let Some(estimate) = DesiredTaskCountHandlers::handle(ev) {
            Ok(nb_ctx.plan_with_task_count(plan, estimate.task_count.limit(nb_ctx.max_tasks()?)))
        } else {
            // We could not determine how many tasks this leaf node should run on, so
            // assume it cannot be distributed and use just 1 task.
            Ok(nb_ctx.plan_with_task_count(plan, Maximum(1)))
        };
    }

The actual property we want is to know is if this assumption was applied:

// We could not determine how many tasks this leaf node should run on, so
// assume it cannot be distributed and use just 1 task.

If it was, that count of 1 would propagate up the stage and we'd never get to this code anyways, right? In which case we could just replace this block with an a plan_err!. Nevermind, this would mean the validator relies on the behavior of the thing it's validating to be correct

[jayshrivastava]

The actual property we want is to know is if this assumption was applied

Still reviewing but the DesiredTaskCountHandler is sort of like a hint right? Because the thing that actually scales up leaves is the ScaleUpLeafNodeHandler. Ideally we inspect the plan directly.

[barbarj]

Still reviewing but the DesiredTaskCountHandler is sort of like a hint right? Because the thing that actually scales up leaves is the ScaleUpLeafNodeHandler. Ideally we inspect the plan directly.

Hmmm. Yeah, but ScaleUpLeafNodeHandler doesn't carry info about intent from what I can tell, just whether or not the scale-up happened.

And actually, really, what we want to know is "is this leaf's data replicated or partitioned", which using the desired task count is just a heuristic for. I wonder if we can encode that some other way.

[gabotechs]

That hardcoded 2... there must be a more robust way of doing this.

@gene-bordegaray, I'd love if you can chime in here.

[barbarj]

PartitioningSatisfaction isn't public in Datafusion 54. It looks like it will be in 55 though.

[barbarj]

The normalize_collect_joins.rs and test from this PR are mostly ready to go, with just some minor cosmetic comments. For shipping those eariler, I'd suggest to split this PR in two:

1. normalize_collect_joins.rs and the tests.
2. validate_stages.rs

I'd like for other people to put some more eyes on 2), but 1) is something we can start shipping sooner. WDYT?

That works for me and lines up with my confidence level about these changes too. I'll cut a separate PR with just the tests and plan fixes. (edit: Done. See #583)

Will review the comments on 2) after that.

[jayshrivastava]

I think this resets all dynamic filters inside the HashJoinExec. I think we need HashJoinExecBuilder::from.

Probably worth having a unit test for collect_left_to_partitioned which ensures no internal fields are reset.

[barbarj]

Oh shoot you're right. Will fix in #583.

[barbarj]

Addressed in #583.

[jayshrivastava]

I think we need to transform full joins right?

In your PR desc:

HashJoin CollectLeft (join types: Left/LeftSemi/LeftAnti/LeftMark/Full) → change
CollectLeft->Partitioned (keys exist; no replication needed)

Right now we exclude full joins from all 3 join layouts (partitioned, broadcast, and single task).

[barbarj]

Yeah, you're right. Removing the gate and fixing the comments in #583

[barbarj]

Addressed in #583.

[jayshrivastava]

The actual property we want is to know is if this assumption was applied

Still reviewing but the DesiredTaskCountHandler is sort of like a hint right? Because the thing that actually scales up leaves is the ScaleUpLeafNodeHandler. Ideally we inspect the plan directly.