feat: store start/end bounds in Partitioning::Range - #26593

👋🏽 @stuhood @saadtajwar @avantgardnerio would any of you be interested in reviewing this? This is related to #22395, which you all have participated in. Please lmk!

👋🏽 @stuhood @saadtajwar @avantgardnerio would any of you be interested in reviewing this? This is related to #22395, which you all have participated in. Please lmk!

Hey @jayshrivastava - thanks for tagging, happy to review this! I'll try to get to this in the next couple of days :)

Thanks @saadtajwar!

FAQ

1. Why make bounds optional instead of using type::MIN and type::MAX

Some types don't have a max or a reasonable value ex. what's the max of a string? "zzzzzz"? For variable length types too

2. Why is this PR so large?

Tests around range satisfaction. Generally, a join may join two different ranges, with or without bounds. 50% of this PR is just adding test coverage around that. For example, the join type influences the behavior here. Happy to delete tests if people thing any are less useful.

Should we settle the alternative @stuhood raised in datafusion-contrib/datafusion-distributed#628 (comment) before bounds become part of the public range contract? There, every task keeps the full Range(N, split_points) and the partitions outside its slice are empty, which needs no new fields on RangePartitioning.

Stu's approach is valid. I also have a similar one in datafusion-contrib/datafusion-distributed#782. I decided to publish this PR because I felt that those two approaches diverge from our goal of keep DFD as close to DF as possible and as thin of a layer over the single node implementation. Introducing these bounds reduces the complexity of datafusion distributed very significantly, and the lines of code as well.

Ballista opted to implement range shuffles without introducing bounds to DF and you can see the tradeoffs that needed to be made.

Consider that we want to distributed this across two tasks

Global: Range(4, splits=[10, 20, 30])

          p0           p1           p2           p3
     [-inf, 10)    [10, 20)     [20, 30)     [30, +inf)
     \___________________/     \____________________/
             Task 0                    Task 1

In ballista, the RangeShuffleReaderExec reports UnknownPartitioning (see reader properties, task slicing)

                         Range shuffle
                         /           \
                        v             v
                Task 0                 Task 1
           RangeShuffleReader     RangeShuffleReader

  Actual ranges:
           [-inf,10) [10,20)      [20,30) [30,+inf)

  DataFusion property:
           UnknownPartitioning(2) UnknownPartitioning(2)

The UnknownPartitioning makes that within a task, we don't have the range partitioning property anymore. Now, a custom plan node won't know its partitioned and may not be able to make certain optimizations.

Your suggestion is to have each task advertise all partitions but mark some as empty like so

                    p0       p1       p2       p3
  Task 0:           data     data     empty    empty
  Task 1:           empty    empty    data     data

  Both advertise Range(4, splits=[10,20,30]).

That is a reasonable alternative, but those empty partitions have an extra cost, the RepartitionExec spawns a reader task and a completion task for every input partition.

Spark describes the same kind of slice by index: a coalesced AQE read gives each reader a span of reducer partitions and keeps the original RangePartitioning with a smaller count (AQEShuffleReadExec).

On Spark, I think the distinction is that its physical RangePartitioning describes ordering and a partition count; it does not contain concrete split values. You can prove that two tasks have disjoint partitions / ranges because the indexes are disjoint. Ie. one task gets 0,1 and another gets 2,3. In datafusion, you can't prove that partitions are disjoint without bounds.

As I read your reply in datafusion-contrib/datafusion-distributed#628 (comment) and the padding thread here, this PR still pads tasks with empty partitions, so what the bounds save is the N - N/W idle partitions on each task. In exchange, every range rule has to account for bounds. For example, equality now includes them, so can_interleave rejects inputs that differ only in bounds, and the layout relationships in #25483 will need to define how bounds compose. Is the idle-partition overhead the remaining cost of that approach, or is there a plan shape it can't express?

Regarding idle partitions, the user has full control here. They can choose to not have them. This PR only introduces the option to pad.

On can_interleave, I would say the current implementation is correct. We would not interleave these:

Input A:  [0,10) [10,20) [20,30) [30,35)
Input B:  [0,10) [10,20) [20,30) [30,40)

However, we may want to interleave them and use this layout.

Output:   [0,10) [10,20) [20,30) [30,40)

Because this PR is a bit large already, I can file a follow up issue for this. Because bounds are new, I doubt many users will run into this. When someone does, that optimization can be implemented fairly easily.

That is a reasonable alternative, but those empty partitions have an extra cost, the RepartitionExec spawns a reader task and a completion task for every input partition.

Thanks, that tells me the difference is run-time cost and DFD code, not a plan the alternative can't express. Since bounds are permanent once they ship, could we measure that cost before settling on them? Each task pays for N - N/W empty partitions: two tasks spawned per input partition in RepartitionExec, plus an empty stream through each operator above the scan. That grows as N per task and N * W across a stage, so a sweep up to the largest shapes DFD expects would show what bounds save. For example, 16 partitions per task on 100 tasks is Range(1600) with 1584 empty partitions on each task. @stuhood also asked in datafusion-contrib/datafusion-distributed#628 (comment) whether a way to declare empty partitions for every Partitioning would remove that cost for Hash too. I don't see that answered yet, and it would avoid a Range-only field.

You can prove that two tasks have disjoint partitions / ranges because the indexes are disjoint. Ie. one task gets 0,1 and another gets 2,3. In datafusion, you can't prove that partitions are disjoint without bounds.

As I read it, the same index argument works in DataFusion under the empty-partition alternative. Every task advertises the same Range(4, [10, 20, 30]), so partition k covers the same keys on every task, and a task that owns 0, 1 has rows only in those partitions. The overlap in the PR description comes from giving each task its own Range(2, [10]) or Range(2, [30]), which drops the global indexes.

I think the more useful comparison with Spark is where it keeps the slice. Its RangePartitioning carries only an ordering and a count. The sampled bounds stay in the exchange's RangePartitioner, and Spark never treats two range layouts as co-partitioned because the bounds are sampled (RangeShuffleSpec). A coalesced read gives each reader a contiguous span of reducer indexes and keeps RangePartitioning with a smaller count (AQEShuffleReadExec). As I read it, that works because a Spark partitioning describes the whole stage, so no task needs to know its own slice.

Spark's coalescing also covers a case DFD will hit: partitions too small to be worth a worker each. AQE packs contiguous reducer partitions up to a target size (ShufflePartitionsUtil). The read keeps RangePartitioning only while every spec is a contiguous, non-overlapping span (isCoalescedRead). Once it splits a skewed partition, it falls back to UnknownPartitioning (L108-L109). In DataFusion terms, coalescing a span drops the split points inside it, so the result is still a plain Range over a subset of the split points. Should DFD coalesce that way before running a range layout on fewer workers? It doesn't remove the per-task slice question, but it shrinks both sides of the trade-off. Under the empty-partition alternative, each task's empty partitions drop from N - N/W to M - M/W for the coalesced count M.

CockroachDB is the other range-partitioned engine I checked. It splits scans into per-node key spans (PartitionSpans), but its optimizer has no range-layout property (physical.Required). Its key-span range router is the closest match to this PR, including the error for rows outside every span. Bulk jobs such as restore use it, not SQL query plans.

So neither system puts a worker's key bounds in a partitioning property that the optimizer compares. DataFusion already goes further than Spark, because RangePartitioning carries split points and drives co-partitioning. Each DFD task also runs its own DataFusion plan, so its partitioning can describe only its own slice. Bounds fill that gap with key values. Empty partitions fill it by giving each task the stage-wide layout that Spark has. Bounds are Range-only, and every rule that compares range layouts (equality, can_interleave, #25483) has to define what they mean. Empty partitions need no new API but cost the partitions above. Could the PR description lay out that trade-off, with the measurement, so we can decide on it?

Because this PR is a bit large already, I can file a follow up issue for this.

A follow-up works for interleaving. As I read it, that's optimizer behavior behind the existing can_interleave signature. Could you file the issue and link it here?

The other half of my question decides what the public API means, so I don't think it can wait for a follow-up: how bounds take part when layouts are compared. Today has_same_layout ignores bounds and == includes them. #25483 will build its relationships on one of those, and its SourceFiner case needs bounds across tasks. For example, the last partition of Range(2, [10], end=20) maps into the first partition of Range(2, [20]) only because of end=20. If #25483 needs bounds to mean something else, we'd be changing the documented semantics of a released type. The API health policy asks us to avoid breaking changes and to deprecate instead, but a change to what == means can't be deprecated, so downstream code would change behavior without warning. Could we settle in this PR how bounds take part in layout comparison and document it on has_same_layout and PartialEq?

the difference is run-time cost and DFD code, not a plan the alternative can't express

I think there's an information difference too. Under the empty-partition alternative, a task advertises Range(16, ...), and which partitions are empty lives only in the source's private state. Nothing in the plan says the task's keys fall within [cut3, cut7), so no rule running in that task can use it. With bounds, the task's partitioning states its own key range.

Bounds aren't used for optimization yet, but they're what pruning the other side of a join, tightening dynamic filters, or key statistics would need. Getting the same information under the alternative would need a way to declare empty partitions on Partitioning, along the lines stuhood suggested, which would be a separate API.

More generally, a property the plan doesn't expose is one no rule can act on, so I'd rather the partitioning describe the task's actual key range than leave it implicit in the source.

Ballista is in the same position as DFD here: each task runs its own DataFusion plan over a slice of the global partitions.

For #25483, I don't anticipate much changing. I think we'll need to add code related to layouts and in that implementation, we can handle bounds, similarly to to can_interleave and compatible_co_partitioning_layout. (I filed #27106 for that by the way to make sure we update those).

For layouts, here's an example where we might want to account for bounds. Suppose a task has Range(2, [10], end=20), and the destination has Range(2, [20]), we would want this behavior:

  Source                            Target

  S0: (-inf, 10) ----+
                    +-------------> T0: (-inf, 20)
  S1: [10, 20) ------+

                                    T1: [20, +inf)
                                        receives nothing

Without bounds, this could happen in theory.

  Source                            Target

  S0: (-inf, 10) ------------------> T0: (-inf, 20)
                                    ^
  S1: [10, +inf) -------------------+
                    |
                    +-------------> T1: [20, +inf)

Since no one has really started on #25483, this is pretty speculative. I think we could find a way to ensure the desired behavior though.

I'm still thinking about has_same_layout and adapt. Will need a bit more time.

@mbutrovich I've thought about partitioning satisfaction and layouts a little bit. I this is the main issue to address before merging.

Background

• symmetric_join_output_partitioning sets the output partitioning of the join based on the join type only (does not care about split points, distribution etc.)
• EnforceDistribution optimizer rule
  • The join declares input_distribution_requirements, so we assert the inputs are key partitioned. However, we don't check split points, bounds, or partition count. We insert a RepartitionExec if this failes
  • enforce_distribution_relationships check that all join inputs are co-partitioned
    • adapt - Imo this should be called adapt_schema and modify the schema temporarily to check co-partitioned ness. I have reverted the previous change were we strip the bounds.
    • has_same_layout - This is the main thing that checks partition count, split points etc. If the layout is not the same, we insert a RepartitionExec

So the question is - now that there's bounds, are there soundness / correctness issues?

Join Example

For the most part, everything in vanilla datafusion is correct. The symmetric_join_output_partitioning does the heavy lifting, choosing the right bounds as needed.

The only suspicious case is this:

HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0)]
  RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1
    DataSourceExec: partitions=1, partition_sizes=[1]
  DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1]

Left partitioning: Range([a@0 ASC], [(10), (20), (30)], 4)
Right partitioning: Range([a@0 ASC], [(10), (20), (30)], 4, start=(0), end=(35))
Output partitioning: Range([a@1 ASC], [(10), (20), (30)], 4, start=(0), end=(35))

We repartition without bounds into a join that has bounds. In this scenario, because the operator is a join, it's expected that the join drops rows outside the bounds. The key thing to note is bounds are not considered by the optimizer rule

• input_distribution_requirements enforces distributions, which are not aware of partition counts, split points, nor bounds
• has_same_layout - considers everything but bounds

The ExecutionPlan must truncate to the bounds.

Custom Plan

Say I create a custom plan which enforces an input distribution, I can introduce an issue where we emit rows outside the bounds. The plan node itself has to override its partitioning to be Range with bounds. EnforceDistribution does not enforce bounds so the plan node emits data outside the bounds.

Output partitioning: Range([a@0 ASC], [(20)], 2, start=(10), end=(30))
Input partitioning: Hash([a@0], 2)

TestProjectionExec // configures the `input_distribution_requirements` to be key partitioned on `a`
  RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2
    DataSourceExec: partitions=2, partition_sizes=[1, 1]

This emits rows outside of [10, 30]. This is a real problem.

Ideas

My two best ideas are as follows:

1. Leave this as is, defer to the operator to handle bounds. If an operator declares Partitioning::Range with bounds and its input doesn't have bounds, the operator should handle this. This is reasonable because the plan node has to opt in. There wouldn't be a correctness issue without that. This is currently true for all joins and works fine.

2. We could make has_same_layout consider bounds, but for joins, the operator handles rows outside the bounds already. Updating has_same_layout to return false when the input has no bounds and the output is bounded would insert repartitions which are unnecessary. For example here, we don't need to repartition the left

HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0)]
    DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1]
    DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1]

Left partitioning: Range([a@0 ASC], [(10), (20), (30)], 4)
Right partitioning: Range([a@0 ASC], [(10), (20), (30)], 4, start=(0), end=(35))
Output partitioning: Range([a@1 ASC], [(10), (20), (30)], 4, start=(0), end=(35))

If I were to make has_same_layout return false on a bounds mismatch, I would add a FilterExec like so just to be safe in lieu of a repartition (an in some cases you would need both a repartition and filter).

FilterExec: 0 <= a AND a < 30

I wouldn't mind putting that in this PR if it makes sense.

3. This one is related to Add Range Layout Replationships to Avoid Repartitions #25483 and is a bigger change. I would say we can do this whenever work on Add Range Layout Replationships to Avoid Repartitions #25483 begins and not put it in this PR.

The solution there would be to probably add partitions to input_distribution_requirements (or RangeLayoutRelationshipKind as the issue proposes).

This change feels like it may be overdue. I'm a new contributor, but I feel like the distributions were added for things like aggregations, where you only care that there's a distinct value in each partition but you don't care about anything else. We're kind of piggybacking on that mechanism for strong partitioning guarantees, like for co-partitioned joins where the actual set of values in each partition matters.

I'm going to push a commit with some tests which show the above examples.

Please let me know what you think!

Re: @stuhood

Landing < to <= split-point validation separately right now

+1

per-partition extrema ... into Statistics

I think Statistics represent observations of the data ("what was the lowest value we've seen"), whereas bounds represent "which values is this partition responsible for", which answers a different question. There may be values that are split points / bounds that do not exist in the data at all.

Ballista's AQE mode doesn't need an unresolved/pending Range partitioning

Because the entire Ballista job is planned as 1 DataFusion plan, then split, unresolved split points would be useful: the job-level plan could declare that it was range partitioned before the split points are resolved and the plan is sliced into task plans (after which it would work as you describe). It's not necessary for that functionality to be in this PR, but until it is available Ballista will have to stick with UnknownPartitioning(K) in such cases.

I pulled the equal split points work in to #27384. Once that's merged, I'll give this one a rebase.