Lazily set the producer head at execution time

[gabotechs]

This is one PR from the following stack of PRs:

• No eager buffering in network connections #477
• Add NetworkBoundaryBuilder argument to inject_network_boundaries #463
• Add MaxGauge metric #464
• Lazily set the producer head at execution time #478 <- you are here
• Refactor task spawner into QueryCoordinator and StageCoordinator #479
• Add cost model #486
• Adaptive task count assignation #432

---

Introduces the ProducerHead type:

pub enum ProducerHead {
    /// No specific head node is necessary.
    None,
    /// The head node should be a [BroadcastExec].
    BroadcastExec { output_partitions: usize },
    /// The head node should be a [RepartitionExec].
    RepartitionExec { partitioning: Partitioning },
}

Which is passed over the network while remotely executing tasks in order to set the appropriate node at the head of a stage.

Today, this is a noop because the right head node in stages is ensured statically at planning time, but in follow up PRs, network boundaries can get swamped and reorganized arbitrarily.

One example that happens in AQE:

1. A JOIN is planned as a CollectLeft

HashJoinExec: mode=CollectLeft
  CoalescePartitionsExec:
    [Stage 1] => NetworkBroadcastExec
      BroadcastExec
        DistributedLeafExec: unknown size
  DistributedLeafExec: unknown size

2. While collecting runtime statistics, it happens that Stage 1 is huge, and during AQE the JOINs are swapped

HashJoinExec: mode=CollectLeft
  DistributedLeafExec: small size
  CoalescePartitionsExec:
    [Stage 1] => NetworkBroadcastExec
      BroadcastExec
        DistributedLeafExec: big size

3. The Stage 1 is now on the probe side, so it needs to be rewritten to a NetworkShuffleExec, otherwise duplicate data will be returned:

HashJoinExec: mode=CollectLeft
  DistributedLeafExec: small size
  CoalescePartitionsExec:
    [Stage 1] => NetworkShuffleExec
      RepartitionExec // <- dynamically swapped at runtime based on the passed `ProducerHead`
        DistributedLeafExec: big size

Passing a ProducerHead at execution time unlocks two things:

1. dynamically set the fanout width accounting for a dynamically scaled upper stage
2. dynamically set the correct operator BroadcastExec or RepartitionExec based on the network boundary above, which might have changed because of AQE

[jayshrivastava]

Why not move this into the match below and error if something unexpected happens? Ex. if you match on ProducerHead::BroadcastExec and input.downcast_ref::<BroadcastExec>() is None, then error?

[gabotechs]

There's not really anything unexpected that can happen here. Actually, the only reason for this function to return a Result<> is because upstream RepartitionExec::try_new() signature requires it.

If the passed ProducerHead is a ProducerHead::BroadcastExec, but the input node is a RepartitionExec, this is fine and expected, it just means that we need to swap the RepartitionExec by a BroadcastExec.

This happens typically while re-ordering JOIN operations on the fly, maybe we swap what was a shuffle join by a broadcast join, and we need to swap the producer head at runtime.

Of course, that example is not possible in this PR, but it will start happening relatively soon.

[jayshrivastava]

mega nit: in upstream they add backticks like [NetworkBoundary]

[gabotechs]

Yeah, I've been noticing for a while. I think this has to do in how things render in docs.rs. We should probably actually:

1. Render this on docs.rs
2. If we think it's indeed weird, do a PR fully dedicated to adding backticks in all docs at once

[gabotechs]

Created this issue for tracking:

#495

[jayshrivastava]

nit: maybe NetworkBoundaryProducer is a better name. I don't feel that strongly though.

[gabotechs]

I like the fact that the "Head" keyword is in there, as it states that we are not referring to a full plan, and we are referring to just the head of it.

[jayshrivastava]

Isn't the ProducerHead of a NetworkShuffleExec always a RepartitionExec with Hash partitioning? So the output partitioning of the NetworkShuffleExec is Hash as well.

[gabotechs]

Today, that happens always yes, but there are legitimate cases for this to be a RoundRobin shuffle: whenever a JOIN order in a CollectLeft join is swapped:

• The left side goes to the right, moving from a broadcast to a round robin shuffle (as we don't have a set of expressions to hash)
• The right side goes to the left, and whatever it was, it gets swapped to being a broadcast

I needed to change this line here today just because otherwise, this benchmark will try to set a ProducerHead in the producing stage to RepartitionExec(UnknownPartitioning), which fails to build in DataFusion. Instead, a RoundRobin is passed so that we can build that RepartitionExec successfully, which does not really matter what it is here, as this is just a benchmark for the transport part.