Skip to content

Commit 7f7d342

Browse files
timsaucerclaude
andcommitted
Cover RoundRobinBatch, stop claiming a plan can report Range
The `scheme` docstring listed four values as though a plan could report any of them. Two of the four needed opposite corrections. `RoundRobinBatch` is reachable, and now tested. It is easy to miss because it never survives at the root: the optimizer inserts one only above a source with fewer partitions than `target_partitions` and CPU work above it to parallelize, so a single-file Parquet scan under a filter and a grouped aggregate produces `RepartitionExec: partitioning=RoundRobinBatch(8)` three levels down. The test walks `children` and asserts all three schemes in that tree, which also pins that the accessor reads each node's own partitioning rather than the root's. `Range` is the other way: it is in the upstream enum, but nothing constructs one in a physical plan. `RangePartitioning` says optimizer and execution support is deliberately unimplemented, per apache/datafusion#22395. Say so, and say why the match arm exists anyway -- it keeps this getter compiling when that support lands, and dropping it would make the match non-exhaustive. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
1 parent 4ab5fc4 commit 7f7d342

3 files changed

Lines changed: 40 additions & 6 deletions

File tree

‎crates/core/src/physical_plan.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,10 +158,16 @@ pub struct PyPhysicalPartitioning {
158158
impl PyPhysicalPartitioning {
159159
/// Which partitioning scheme this is.
160160
///
161-
/// One of `RoundRobinBatch`, `Hash`, `Range`, or `UnknownPartitioning`.
162161
/// `UnknownPartitioning` is what a plan reports when it knows how many
163162
/// partitions it has but nothing about how rows are distributed between
164-
/// them, which is the common case for a file scan.
163+
/// them, which is the common case for a file scan. `RoundRobinBatch` and
164+
/// `Hash` come from a `RepartitionExec`.
165+
///
166+
/// `Range` is in the upstream enum but no plan reports it yet: optimizer
167+
/// and execution support is deliberately unimplemented, per
168+
/// <https://github.com/apache/datafusion/issues/22395>. The arm is here so
169+
/// this getter keeps compiling when that lands, not because it is
170+
/// reachable today.
165171
#[getter]
166172
pub fn scheme(&self) -> &'static str {
167173
match self.partitioning {

‎python/datafusion/plan.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -346,10 +346,11 @@ def __init__(self, partitioning: df_internal.PhysicalPartitioning) -> None:
346346
def scheme(self) -> str:
347347
"""Which partitioning scheme this is.
348348
349-
One of ``"RoundRobinBatch"``, ``"Hash"``, ``"Range"``, or
350-
``"UnknownPartitioning"``. A plan reports ``"UnknownPartitioning"``
351-
when it knows how many partitions it has but nothing about how rows
352-
are distributed between them, which is the usual case for a file scan.
349+
``"UnknownPartitioning"`` means the plan knows how many partitions it
350+
has but nothing about how rows are distributed between them, which is
351+
the usual case for a file scan. ``"RoundRobinBatch"`` and ``"Hash"``
352+
come from a repartition the optimizer inserted. ``"Range"`` is defined
353+
upstream but no plan reports it yet.
353354
354355
Examples:
355356
>>> from datafusion import SessionContext

‎python/tests/test_plans.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import datetime
1919

2020
import pyarrow as pa
21+
import pyarrow.parquet as pq
2122
import pytest
2223
from datafusion import (
2324
ExecutionPlan,
@@ -144,6 +145,32 @@ def test_output_partitioning_reports_the_scheme_not_just_the_count() -> None:
144145
assert repr(partitioning) == "Hash([a@0], 4)"
145146

146147

148+
def test_output_partitioning_reports_round_robin(tmp_path) -> None:
149+
"""A round-robin repartition reports `RoundRobinBatch`.
150+
151+
The optimizer only inserts one above a source with fewer partitions than
152+
`target_partitions` and CPU work above it to parallelize, and it never
153+
survives at the root, so reach it by walking `children`.
154+
"""
155+
path = tmp_path / "rr.parquet"
156+
pq.write_table(pa.table({"a": list(range(2000)), "b": [1] * 2000}), path)
157+
158+
ctx = SessionContext(SessionConfig().with_target_partitions(8))
159+
ctx.register_parquet("t", str(path))
160+
plan = ctx.sql("select a, sum(b) from t where a > 5 group by a").execution_plan()
161+
162+
schemes = set()
163+
stack = [plan]
164+
while stack:
165+
node = stack.pop()
166+
schemes.add(node.output_partitioning.scheme)
167+
stack.extend(node.children())
168+
169+
# The single-file scan, the round-robin above it, and the hash repartition
170+
# for the grouping are all in one tree.
171+
assert schemes == {"UnknownPartitioning", "RoundRobinBatch", "Hash"}
172+
173+
147174
def test_execute_rejects_an_out_of_range_partition() -> None:
148175
"""An out-of-range partition index raises instead of panicking."""
149176
ctx = SessionContext()

0 commit comments

Comments
 (0)