Preserve range partitioning across network boundaries

[mdashti]

Summary

Preserve Partitioning::Range across stage boundaries so downstream consumer operations (joins, aggregations, window functions) can remain co-partitioned without triggering unnecessary hash repartitions.

Background

The original description proposed unioning matching partition indices across tasks in NetworkCoalesceExec (output_partition[i] = Union(task[*].partition[i])). Topologically, that multi-producer exchange is a shuffle rather than a coalesce.

With datafusion-contrib/datafusion-distributed#730 adding range support to NetworkShuffleExec and apache/datafusion#25483 introducing range layout relationships upstream, the problem separates into two distinct boundaries:

1. NetworkShuffleExec (Sparse Exchange): When multiple producer tasks output the same range buckets, we need sparse connections so consumer i only reads bucket i across producers without per-row repartitioning.
2. NetworkCoalesceExec (Disjoint Gather): When producer tasks already hold disjoint, contiguous ranges (e.g. leaf file groups), coalescing them into fewer consumer tasks should yield a coarser SourceFiner RangePartitioning. Currently, it degrades to UnknownPartitioning:

datafusion-distributed/src/execution_plans/common.rs

Lines 30 to 38 in 497b5b3

	// serve, so drop to an unknown layout with the scaled count; the boundary math
	// and the consumer run on counts alone.
	// TODO(#68): carry the range property across the boundary so consumer-side
	// joins can stay co-partitioned.
	Partitioning::Range(range) => Ok(Partitioning::UnknownPartitioning(f(
	range.partition_count()
	))),
	}
	}

---

Impacted Plan Shapes

Queries that execute a co-partitioned operation, reduce rows, and feed into a downstream stage expecting the same range key:

1. Multi-way join with early row reduction

SELECT *
FROM (
    SELECT t1.id, t1.val, t2.extra
    FROM t1 JOIN t2 ON t1.id = t2.id
    WHERE t1.val > 100
) sub
JOIN t3 ON sub.id = t3.id
ORDER BY sub.id
LIMIT 10;

• Stage 1: Co-partitioned join of t1 and t2 on id followed by a selective filter.
• Stage 2: Directly join reduced stream with t3 (also range-partitioned on id) without reshuffling.
• Current issue: NetworkCoalesceExec drops range partitioning to UnknownPartitioning, while NetworkShuffleExec cannot prove co-partitioning if split points or partition counts differ from t3, forcing a full hash repartition.

2. Group-by aggregation or window function over co-partitioned join

SELECT t1.category_id, count(*), sum(t2.amount)
FROM t1 JOIN t2 ON t1.category_id = t2.category_id
GROUP BY t1.category_id;

• Workers execute a co-partitioned join and partial aggregate.
• Losing range layout across the stage boundary forces an intermediate hash shuffle instead of letting the final aggregate/window run in place per partition.

---

Action Items

1. Leverage upstream Add Range Layout Replationships to Avoid Repartitions apache/datafusion#25483 to enable sparse NetworkShuffleExec connections and whole-batch forwarding between matching range partitions.
2. Enable downstream distribution enforcement to accept compatible SourceFiner / TargetFiner range layouts for joins and aggregations.
3. Update NetworkCoalesceExec to merge split points when gathering contiguous, disjoint task ranges rather than degrading to UnknownPartitioning.

[stuhood]

Thanks a lot! We'll want to upstream #67 and this fix atop datafusion-55.