Need for Customizable Plan Annotation and Network Boundary Logic in DFD

[gabrielkerr]

Summary:
This issue aims to start a discussion around improving extensibility in datafusion-distributed, especially for custom plan annotations and network boundaries. I would appreciate insights from the DataFusion community on potential design directions and best practices.

---

Challenge and Motivation

I believe there is significant value in expanding the extensibility of datafusion-distributed (DFD). The project’s core strengths—plan annotation, insertion of network boundaries, and distribution of sub‑plans to workers—make it a natural place for more flexible customization.

My colleagues and I have been working toward implementing custom network boundaries and plan annotations in a fork of DFD. The use case involves inserting multiple ExecutionPlan nodes instead of relying solely on NetworkShuffleExec, NetworkCoalesceExec, or NetworkBroadcastExec. In practice, this requires a mechanism to introduce custom plan annotations and network boundaries beyond what DFD currently supports.

An initial attempt at introducing this extensibility can be found in this draft PR by @kurtvolmar:
kurtvolmar#1.

However, given that DataFusion itself already provides many extension points, it may not be ideal for DFD to introduce a separate, parallel extensibility framework. Instead, it seems likely that a DataFusion‑native mechanism could be leveraged to support more flexible behavior within DFD.

---

Proposal and Call for Collaboration

The existing DistributedPhysicalOptimizerRule encapsulates a substantial amount of logic. One possible direction would be to decompose this rule into smaller, more focused components—such as a PlanAnnotationRule and NetworkBoundaryRule—and expose configuration or hooks that allow users to implement custom logic where needed.

Community input would be highly valuable, particularly around:

• Whether splitting DistributedPhysicalOptimizerRule into smaller, pluggable rules aligns with the project’s direction.
• Alternative approaches in DataFusion that could enable the desired extensibility without modifying DFD directly.
• Prior art or patterns in the DataFusion ecosystem that could help inform a clean design.

Feedback, suggestions, or discussion from maintainers and contributors would be greatly appreciated. The goal is to collaborate on a design that increases flexibility without adding unnecessary complexity to DFD.

cc: @gabotechs

[gabotechs]

Splitting the DistributedPhysicalOptimizerRule into smaller rules sounds good 👍, from an extensibility and maintainability standpoint.

There's one challenge that comes to mind: The current DistributedPhysicalOptimizerRule keeps a record of the original plan, so that if the plan doesn't end up getting distributed, we returned the original one. By separating into multiple rules, we'll lose the original plan along the first couple of rules, and that would make the last rule unavailable of returning the original one, as it got lost.

At minimum I can think of a way of workarounding this, so we have options.

I'll draft a PR for this.

[kurtvolmar]

@gabotechs based the open Issue apache/datafusion#20396 you created, and previous discussions, is the approach you are considering a chain of annotation optimizers, followed by an optimizer which performs the plan mutation? (i.e. something likeNetworkShuffleAnnotationOptimizer, NetworkCoalesceAnnotationOptimizer, NetworkBroadcastAnnotationOptimizer, PlanDistributionOptimizer)

If so, I would like to consider an alternative approach if annotations are not added to the ExecutionPlan. Instead of having separate annotation and mutation optimizers, what if each optimizer were responsible for it's own plan mutation and task estimation, producing a correct distributed plan as output. Subsequent optimizers would need to recalculate tasks. This approach may be limiting for some of the task calculation and partition scaling, but may be worth considering as an approach.

[gabotechs]

is the approach you are considering a chain of annotation optimizers, followed by an optimizer which performs the plan mutation?

Not really, that's not what I had in mind.

Instead of having separate annotation and mutation optimizers, what if each optimizer were responsible for it's own plan mutation and task estimation, producing a correct distributed plan as output

Producing correct distributed plans requires mutations in nodes like BroadcastExec or RepartitionExec that are hard to rollback in order to adequate them to different task configurations. I think this can get complicated very soon.

[gabotechs]

About apache/datafusion#20396, if that does not end up moving forward, I have a backup plan for still providing users the ability to annotate their plans, maintaining the original idea but with a not-so-nice API, at least not as nice as having it integrated in DataFusion ExecutionPlans itself.

Just hold tight, I'll have something ready early next week

[gabotechs]

I came up with this #357, let me know what you guys think.

This test demonstrates how to customize a distributed plan with the new API:

https://github.com/datafusion-contrib/datafusion-distributed/blob/2dd97539bf736e83c1697b763c5ba1c724a1e69d/tests/custom_network_boundary_placeholders.rs

[gabotechs]

Thanks for the feedback on #357 guys!

Before jumping into more code, it would be nice to discuss here what parts of #357 make sense to have and what other parts are better addressed in a different way.

I agree that, in the name of modeling things in terms of existing DataFusion extension points, #357 might go a bit too far with things that I'd consider "workarounds", like the DistributedContext thing.

Regardless of the rule division, what do you think about the NetworkBoundaryPlaceholder as a public facing API for injecting custom network boundaries?

[kurtvolmar]

Thank you @gabotechs! We are planning to revisit this tomorrow to continue the discussion.

[gabrielkerr]

Hey @gabotechs thanks for the comment replies. To answer your question, I think the NetworkBoundaryPlaceholder as a public facing API is nearly there. Can we compare and contrast this API with NetworkBoundaryStrategy in kurtvolmar#1?

I like that you've composed the DistributedPhysicalOptimizerRule into smaller rules and guardrail rules like ensuring the plan has a single coalesce at the top. I'm afraid that there is a lot of added complexity with how granular it has been broken down. The NetworkBoundaryStrategy that @kurtvolmar has implemented buys us essentially the same thing, but the drawback is that it's not implemented as a PhysicalOptimzerRule and introduces an extension point that differs from what Datafusion offers.

What if we refactored @kurtvolmar's approach to use PhysicalOptimizerRules and similar guardrails to what you have done in #357? This may marry the simplicity of our initial approach with the Datafusion "friendliness" of yours. We're happy to give this a shot, but first what are your thoughts?

[kurtvolmar]

@gabotechs really like the composability and general approach of your PR. I think you narrowed down the key tension very well, being the NetworkBoundaryPlaceholder API. I think it's a solid first pass, but some tweaks to the approach may make it a bit more expressive. I'll add to the feedback from @gabrielkerr above.

For the NetworkBoundaryPlaceholder API my primary concern is that within the proposed composition of optimzers, it doesn't simplify the user implementation for distributed planning as much as I'd like to see. In our use case of implementing a distributed Tree Reduction plan, the implementation with a NetworkBoundaryPlaceholder and just inserting NetworkCoalesceExec look very similar, but the latter doesn't require the concept of NetworkBoundaryPlaceholder:

With NetworkBoundaryPlaceholder

• Create a TreeReductionOptimizer which:
  • traverses the plan and identifies the a hash repartiton
  • recursively inserts a hierarchy of NetworkBoundaryPlaceholder with NetworkBoundaryKind::Coalesce along with partial aggregates and partition coalesces.
  • calculates input_task_counts for each stage
• Since both tree reduction and shuffles both trigger based on hash repartitions, the TreeReductionOptimizerand InjectNetworkBoundaryPlaceholders (default optimizer for network boundaries) must be made mutually exclusive. But this would in theory place nicely with ApplyNetworkBoundaries.

This creates a plan structured something like:

ROOT
  ...
  NetworkBoundaryPlaceholder (Coalesce) input_task_count=2
    AggregateExec (Partial)
      NetworkBoundaryPlaceholder (Coalesce) input_task_count=4
        AggregateExec (Partial)
          NetworkBoundaryPlaceholder (Coalesce) input_task_count=16
            ...
             LEAF

With NetworkCoalesceExec

• Create a TreeReductionOptimizer which:
  • traverses the plan and identifies the a hash repartiton (same as above)
  • recursively inserts a hierarchy of NetworkCoalesceExec along with partial aggregates and partition coalesces. (same as above, but inserting NetworkCoalesceExec instead of NetworkBoundaryPlaceholder.
  • calculates input_tasks for each stage (same as above)
• Since both tree reduction and shuffles both trigger based on hash repartitions, the TreeReductionOptimizerand InjectNetworkBoundaryPlaceholders (default optimizer for network boundaries) must be made mutually exclusive. Now that we're just inserting NetworkCoalesceExec directly it is also unnecessary to use ApplyNetworkBoundaries.

This creates a plan structured something like:

ROOT
  ...
  NetworkBoundaryPlaceholder (Coalesce) input_tasks=2
    AggregateExec (Partial)
      NetworkBoundaryPlaceholder (Coalesce) input_tasks=4
        AggregateExec (Partial)
          NetworkBoundaryPlaceholder (Coalesce) input_tasks=16
            ...
             LEAF

---

Both of these approaches are made possible due to breaking up the DistributedPhysicalOptimizerRule into multiple rules, which we definitely want. The shortcomings of the NetworkBoundaryPlaceholder is that, for the above plan, it is almost no different to just directly insert the network boundaries that you want.

The key tension is that custom distributed PhysicalOptimzerRules will likely conflict with the the InjectNetworkBoundaryPlaceholders and ApplyNetworkBoundaries. I would really like to find a solution where custom rules can play nicely with the default network boundaries, rather than having an all-or-nothing approach. To make the NetworkBoundaryPlaceholder API more useful, I think we need to further break up the InjectNetworkBoundaryPlacholders rule into multiple optimizers such as ShuffleOptimzerRule, BroadcastOptimzerRule, etc. This would enable implementations to have better composability for their optimizer rules.

---

Output Tasks
Another addition to NetworkBoundaryPlaceholder which would be extremely useful is an optional output_task_count field which specifies how many output tasks that placeholder will produce. This would give users control into the tasking calculations when needed. This would prevent the automatic scaling that the plan distribution performs, but in the case of Tree Reduction it is extremely important that the top NetworkCoalesceExec reduces to 1 output task.

[gabotechs]

I would really like to find a solution where custom rules can play nicely with the default network boundaries, rather than having an all-or-nothing approach. To make the NetworkBoundaryPlaceholder API more useful, I think we need to further break up the InjectNetworkBoundaryPlacholders rule into multiple optimizers such as ShuffleOptimzerRule, BroadcastOptimzerRule, etc.

🤔 this sounds interesting. It looks a bit complicated, but it might be a good idea.

Another addition to NetworkBoundaryPlaceholder which would be extremely useful is an optional output_task_count field

Note that in this project, there is no concept of output_task_count. Each network boundary decides the number of task of the stage below, but it has no capability of "outputing" to tasks, a network boundary just receives data, but doesn't give. The number of tasks in which a network boundary runs on is decided by the input_task_count of the network boundary above.

traverses the plan and identifies the a hash repartiton (same as above)

If you indeed want to rearrange the plan removing hash repartitions, inserting partial reduce aggregations, etc, it might be a better fit to do this in a different optimizer rule the precedes Distributed DataFusion. Note that there's one important assumption in the plan annotator: the plan should not be modified there, it should just be annotated with network boundary information. You can read more about this design in the PR that introduced it #259.

This is the main reason for us to do all the preparatory steps before reaching the annotations.

---

One issue with kurtvolmar#1 is the amount of API surface impact it has for providing extensibility, and the fact that it wires up at the plan annotator point. I don't think we can afford to make any guarantees about the presence of the current plan annotator, as it's something very susceptible to disappear at any moment.

These are all the new things that would become part of the public API in that PR:

pub struct CombinedNetworkBoundaryStrategy {
}
pub struct NetworkBoundaryAnnotation {
    pub required_network_boundary: Option<PlanOrNetworkBoundary>,
    pub output_tasks: Option<usize>,
}
pub enum PlanOrNetworkBoundary {
    Plan(Arc<dyn ExecutionPlan>),
    Shuffle,
    Coalesce,
    Broadcast,
    Extension(&'static str),
}
pub struct NetworkBoundaryContext<'a> {
    pub boundary_type: &'a PlanOrNetworkBoundary,
    pub new_children: Arc<dyn ExecutionPlan>,
    pub query_id: Uuid,
    pub stage_id: usize,
    pub task_count: usize,
    pub input_task_count: usize,
    pub config: &'a ConfigOptions,
}
pub trait NetworkBoundaryStrategy: Debug + Send + Sync {
    fn annotate_network_boundary(
        &self,
        plan: &dyn ExecutionPlan,
    ) -> Option<NetworkBoundaryAnnotation>;
    fn apply_boundary(
        &self,
        context: &NetworkBoundaryContext<'_>,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
}
pub struct DoGet {
    pub plan_proto: Bytes,
    pub target_task_index: u64,
    pub target_task_count: u64,
    pub target_partition_start: u64,
    pub target_partition_end: u64,
    pub stage_key: Option<StageKey>,
}
pub enum AppMetadata {
    MetricsCollection(MetricsCollection),
}
pub struct DistributedCodec;
pub struct FlightAppMetadata {
    pub partition: u64,
    pub created_timestamp_unix_nanos: u64,
    pub content: Option<AppMetadata>,
}

It seems like the surface area of the API is pretty big, and I think the only thing that should really be needed is this:

pub enum NetworkBoundaryKind {
    Shuffle,
    Coalesce,
    Broadcast,
}
pub struct NetworkBoundaryPlaceholder {
    pub kind: NetworkBoundaryKind,
    pub input_task_count: usize,
    pub input: Arc<dyn ExecutionPlan>,
}

[gabotechs]

A good exercise that comes to mind would be to try to make your custom plan executable first using normal DataFusion:

• Remove any optimizer rule that you'd prefer to not be executed, maybe you'd prefer this ones to not be executed? https://github.com/apache/datafusion/blob/35749607f585b3bf25b66b7d2289c56c18d03e4f/datafusion/physical-optimizer/src/optimizer.rs#L116-L118.
• Add any optimizer rule that you find appropriate that lands your plan in an executable state using single node. This is the place to deal with nodes unrelated to Distributed DataFusion, like the partial reduce aggregation.
• At this point, you should have a plan that is executable with normal DataFusion, which might even be able to demonstrate speed ups.
• Add another rule that injects the network boundaries in the appropriate places. At this point, the resulting plan should be virtually equivalent to the single-node one, but just happens to have network boundaries in between.

Getting figured out first the plan with vanilla DataFusion tools will provide you very good insights about what's really needed in this project from an extensibility stand point.

[kurtvolmar]

@gabotechs I've implemented our custom PhysicalOptimizerRule based on your PR branch. I haven't yet tried out your suggestion to try this out single node with more standard DF optimizers - I plan to do this tomorrow.

I've implemented tree reduction as a PhysicalOptimizerRule in two ways. Both implementations are nearly identical and produced correct plans for us:

1. Directly insert NetworkCoalesceExec with calculated tasks and bypass both InjectNetworkBoundaryPlaceholdersandApplyNetworkBoundaries`. Here is the chain of optimizers:
   • StartDistributedContext
   • AddCoalesceOnTop
   • InsertBroadcast
   • TreeReductionRule
   • InjectNetworkBoundaryPlaceholders
   • ApplyNetworkBoundaries
   • BatchCoalesceBelowBoundaries
   • EndDistributedContext
2. Insert NetworkBoundaryPlaceholder with NetworkBoundaryKind::Coalesce and use along with ApplyNetworkBoundaries, letting ApplyNetworkBoundaries insert the actual NetworkCoalesceExec.
   • StartDistributedContext
   • AddCoalesceOnTop
   • InsertBroadcast
   • TreeReductionPlaceholderRule
   • InjectNetworkBoundaryPlaceholders
   • ApplyNetworkBoundaries (EDIT: this rule is included in planning)
   • BatchCoalesceBelowBoundaries
   • EndDistributedContext

I don't think NetworkBoundaryPlaceholder would be necessary for us to implement and we would likely go with approach 1. I still think the NetworkBoundaryPlaceholder API would be more valuable if InjectNetworkBoundaryPlaceholders were split into separate rules like InjectShufflePlaceholders and InjectBroadcastPlaceholders. This is because the TreeReductionPlaceholderRule is mutually exclusive with InjectNetworkBoundaryPlaceholders specifically for the shuffle. That may be unnecessary if we can do the correct optimizations before it goes into the datafusion-distributed optimizers (again, I'll try this out soon).

[gabotechs]

API would be more valuable if InjectNetworkBoundaryPlaceholders were split into separate rules like InjectShufflePlaceholders and InjectBroadcastPlaceholders

🤔 this is a good idea, it looks complicated but it might be worth trying. One challenge that comes to mind is that a Shuffle is not really something users can not do if there is a RepartitionExec. If DataFusion decided to place a RepartionExec in a plan, it necessarily means that there should be a Shuffle right above, otherwise incorrect results will happen.

This means that probably a good way of avoiding a Shuffle, is to tweak the plan so that there are no RepartitionExecs to begin with, that would be a good and strictly correct way of avoiding Shuffles.

I don't think NetworkBoundaryPlaceholder would be necessary for us to implement and we would likely go with approach 1

This would simplify things definitely

I haven't yet tried out your suggestion to try this out single node with more standard DF optimizers - I plan to do this tomorrow.

👍 Cool, I think there is a good chance this provides you some good hints about next steps.