Allow custom physical operators to opt into NetworkCoalesce boundary placement

Background

datafusion-distributed starts from a normal single-node DataFusion physical plan, then inserts network boundaries where the distributed plan needs to split into stages.

For the final/root collection case, the planner currently adds a CoalescePartitionsExec when the physical plan has multiple output partitions:

let mut plan = Arc::clone(&original_plan);

if plan.output_partitioning().partition_count() > 1 {
    plan = Arc::new(CoalescePartitionsExec::new(plan));
}

plan = insert_broadcast_execs(plan, cfg)?;
plan = inject_network_boundaries(plan, CardinalityBasedNetworkBoundaryBuilder, cfg).await?;

https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/distributed_planner/distributed_query_planner.rs#L101-L111

Then, during boundary injection, NetworkCoalesceExec is inserted when the current node's parent is exactly CoalescePartitionsExec or SortPreservingMergeExec:

} else if let Some(parent) = parent
    && !plan.children().is_empty()
    && (parent.is::<CoalescePartitionsExec>()
        || parent.is::<SortPreservingMergeExec>())
{
    let input_stage = LocalStage {
        query_id: nb_ctx.query_id,
        num: nb_ctx.fetch_add_stage_id(),
        plan: nb_ctx.plan_with_task_count(plan, task_count),
        tasks: task_count.as_usize(),
    };

    let result = nb_ctx
        .nb_builder
        .build(input_stage, TypeId::of::<NetworkCoalesceExec>(), nb_ctx)
        .await?;

    let nb = Arc::new(NetworkCoalesceExec::from_stage(
        result.input_stage,
        result.input_properties,
        1,
    ));

    Ok(nb_ctx.plan_with_task_count(nb, result.consumer_task_count))
}

https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/distributed_planner/inject_network_boundaries.rs#L308-L360

That maps cleanly to the documented role of NetworkCoalesceExec:

/// This is the equivalent of a [CoalescePartitionsExec] but coalescing tasks across the network
/// between distributed stages.

https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/execution_plans/network_coalesce.rs#L21-L25

This works well for built-in DataFusion nodes. The gap is custom ExecutionPlan implementations that have the same distributed boundary requirement, but are not literally CoalescePartitionsExec or SortPreservingMergeExec.

Problem

Some custom physical operators need their child input to cross a coalesce network boundary, but the current automatic planner cannot recognize them.

For example:

CustomCollectExec          output partitions = 1
  ParallelInputExec        output partitions = N

A distributed execution should be able to run the lower part in parallel, then collect the producer-task output through the same NetworkCoalesceExec machinery already used for built-in coalescing nodes.

Conceptually:

CustomCollectExec
  NetworkCoalesceExec
    ParallelInputExec

Today, downstream code has to make this visible to the planner indirectly. The common workaround is to insert a marker CoalescePartitionsExec or SortPreservingMergeExec around the custom node purely so this check matches:

parent.is::<CoalescePartitionsExec>() || parent.is::<SortPreservingMergeExec>()

That works, but it is a bit awkward: the extra node is not there because the query logically needs another DataFusion coalesce operator; it is there because datafusion-distributed currently uses concrete parent types as the signal for placing NetworkCoalesceExec.

Manual NetworkCoalesceExec insertion is another option, but then custom integrations need to take over more distributed-planner detail than necessary.

Proposal

Add a small opt-in for custom operators that need the existing NetworkCoalesceExec boundary behavior.

One possible shape:

pub struct NetworkCoalesceBoundaryContext<'a> {
    pub parent: Option<&'a Arc<dyn ExecutionPlan>>,
    pub child_index: Option<usize>,
    pub producer: &'a Arc<dyn ExecutionPlan>,
}

pub trait NetworkCoalesceBoundaryRule {
    fn requires_network_coalesce(
        &self,
        ctx: NetworkCoalesceBoundaryContext<'_>,
        cfg: &ConfigOptions,
    ) -> Result<bool> {
        Ok(false)
    }
}

The built-in behavior would stay exactly as it is today:

parent.is::<CoalescePartitionsExec>() || parent.is::<SortPreservingMergeExec>()

Registered custom rules would only add more cases where the planner should use that same existing path.

For example:

fn requires_network_coalesce(
    &self,
    ctx: NetworkCoalesceBoundaryContext<'_>,
    _cfg: &ConfigOptions,
) -> Result<bool> {
    Ok(ctx.producer.is::<CustomCollectExec>())
}

The rule would not construct NetworkCoalesceExec itself. It would just tell the planner that this edge should be treated like the existing CoalescePartitionsExec / SortPreservingMergeExec case. Task-count reconciliation, stage construction, boundary preparation, and boundary elision would remain owned by datafusion-distributed.

I am intentionally only suggesting this for coalesce boundaries. Shuffle and broadcast have more specific partitioning/join semantics, and I do not think they need to be part of this change.

Expected behavior

With no custom rule registered, plans behave exactly as they do today.

With a custom coalesce rule registered, a custom physical operator can trigger the same NetworkCoalesceExec insertion currently reserved for built-in coalesce-like parents.

The opt-in contract should be documented clearly: returning true means the custom operator is compatible with collecting producer-task output through NetworkCoalesceExec.

Tests

I think the useful coverage would be:

• existing CoalescePartitionsExec and SortPreservingMergeExec behavior stays unchanged;
• a custom physical operator does not get a coalesce boundary by default;
• the same custom operator gets a NetworkCoalesceExec when a custom rule opts it in;
• root multi-partition collection behavior stays unchanged;
• shuffle and broadcast plans are unaffected.

Start with distributed_planner/inject_network_boundaries.rs, especially the existing CoalescePartitionsExec and SortPreservingMergeExec parent check, then review distributed_query_planner.rs and execution_plans/network_coalesce.rs. Define the opt-in contract and registration shape from the proposal before implementing it. Done means built-in behavior is unchanged, an opted-in custom operator receives NetworkCoalesceExec, unregistered operators do not, and shuffle and broadcast behavior remains unaffected.