Operators declare what they need: how DataFusion's planner inserts the shuffle for you

TL;DR: An operator never builds a shuffle itself. It declares a Distribution requirement (Unspecified, SinglePartition, or KeyPartitioned(exprs)), and the EnsureRequirements physical optimizer rule walks the plan bottom-up, checks whether the child’s current Partitioning satisfies the requirement, and inserts a RepartitionExec only when it doesn’t. Joins add a second, harder constraint: both sides must land in the same partition for the same key.

The problem: who decides where data goes?

Consider SELECT ... FROM a JOIN b ON a.k = b.k. A partitioned hash join wants the following invariant:

text

left:  Hash([k])              right: Hash([k])
  P0: k ∈ {1, 4, 9}              P0: k ∈ {1, 4, 9}
  P1: k ∈ {2, 5}                 P1: k ∈ {2, 5}

Partition i of the left side may only be joined with partition i of the right side. If the scan of a produces 8 file partitions and the scan of b produces 3, somebody has to move rows until both sides agree on the layout. Where should that decision live?

The naive answer — “in the planner, with hard-coded rules per operator type” — is what most query engines do, and it rots. Every new operator needs its rule; every rule needs to know about every other operator’s behaviour; the join rule has to know that the aggregate below it also repartitions. It’s a web of special cases.

DataFusion inverted this. The operator states what it needs; the optimizer satisfies it generically.

The vocabulary: `Distribution` vs `Partitioning`

Two enums, and the difference matters.

Partitioning describes what a node actually produces:

rust
// datafusion/physical-expr/src/partitioning.rs:122-132
pub enum Partitioning {
    RoundRobinBatch(usize),
    Hash(Vec<Arc<dyn PhysicalExpr>>, usize),
    Range(RangePartitioning),
    UnknownPartitioning(usize),
}

Distribution describes what a node requires of its input:

rust
// datafusion/physical-expr/src/partitioning.rs:934-946
pub enum Distribution {
    /// Unspecified distribution
    UnspecifiedDistribution,
    /// A single partition is required
    SinglePartition,
    /// Deprecated historical name for [`Distribution::KeyPartitioned`].
    #[deprecated(since = "55.0.0", note = "Use Distribution::KeyPartitioned")]
    HashPartitioned(Vec<Arc<dyn PhysicalExpr>>),
    /// Requires children to be distributed in such a way that the same
    /// values of the keys end up in the same partition
    KeyPartitioned(Vec<Arc<dyn PhysicalExpr>>),
}

HashPartitioned is deprecated in favour of KeyPartitioned — a naming change that reflects a semantic broadening: the requirement is about key locality, not about the literal hash function. The requirement doesn’t say “partition by hash(k) % 8”; it says “rows with the same k must be in the same partition, and it must be the same partition across children”.

The matching side of this is PartitioningSatisfaction:

rust
// datafusion/physical-expr/src/partitioning.rs:535-554
pub enum PartitioningSatisfaction {
    /// The partitioning does not satisfy the distribution requirement
    NotSatisfied,
    /// The partitioning exactly matches the distribution requirement
    Exact,
    /// The partitioning satisfies the distribution requirement via subset logic
    Subset,
}

Subset is the interesting one. If a child is already hash-partitioned on (k) and the parent needs (k, j), then (k) is not sufficient (rows with the same k but different j would be split) — that’s NotSatisfied. But if the child is hash-partitioned on (k, j) and the parent only needs (k), the child’s partitioning is stronger than required, so it’s satisfied. The rule is implemented in key_satisfaction (partitioning.rs:649-684), and it normalizes both sides through the plan’s equivalence classes first, so a and b count as the same key if the optimizer knows they’re equal.

Declaring the requirement

The trait method is on ExecutionPlan itself:

rust
// datafusion/physical-plan/src/execution_plan.rs:188-210
/// Specifies simple per-child input distribution requirements.
///
/// Deprecated: override [`Self::input_distribution_requirements`] instead.
///
/// By default, each child has [`Distribution::UnspecifiedDistribution`].
#[deprecated(since = "55.0.0", note = "Use input_distribution_requirements")]
fn required_input_distribution(&self) -> Vec<Distribution> {
    vec![Distribution::UnspecifiedDistribution; self.children().len()]
}

/// Specifies the input distribution requirements for this plan.
///
/// The default implementation wraps [`Self::required_input_distribution`].
/// Override this method for richer requirements, such as allowing alternate
/// satisfaction policies or requiring multiple children to be co-partitioned.
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
    InputDistributionRequirements::new(self.required_input_distribution())
}

So the modern API is input_distribution_requirements, and the old required_input_distribution is a compatibility shim for out-of-tree operators. The richer return type exists because per-child independence is not always enough:

rust
// datafusion/physical-plan/src/distribution_requirements.rs:59-66
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct InputDistributionRequirements {
    /// Per-child distribution requirements, indexed by child position.
    children: Vec<ChildDistributionRequirement>,
    /// Child indexes that must also have compatible partition layouts.
    co_partitioned: Option<Vec<usize>>,
}

co_partitioned is the join constraint. InputDistributionRequirements::co_partitioned documents it precisely:

rust
// datafusion/physical-plan/src/distribution_requirements.rs:107-132
/// Create a requirement that all children are co-partitioned.
///
/// ```text
/// left:  Range(left.a ASC,  split_points=[10, 20])
/// right: Range(right.x ASC, split_points=[10, 20])
///
/// partition 0 from both sides contains keys before 10
/// partition 1 from both sides contains keys in [10, 20)
/// partition 2 from both sides contains keys at/after 20
/// ```
///
/// If the split points differ, partition `i` from one side no longer covers
/// the same key range as partition `i` from the other side.

Who declares what

A partial aggregate wants any distribution — it will merge later:

rust
// datafusion/physical-plan/src/aggregates/mod.rs:2514-2526
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
    InputDistributionRequirements::new(match &self.mode {
        AggregateMode::Partial | AggregateMode::PartialReduce => {
            vec![Distribution::UnspecifiedDistribution]
        }
        AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned => {
            vec![Distribution::KeyPartitioned(self.group_by().input_exprs())]
        }
        AggregateMode::Final | AggregateMode::Single => {
            vec![Distribution::SinglePartition]
        }
    })
}

A window function wants its PARTITION BY keys co-located — and if it has no PARTITION BY at all, it wants everything in one place:

rust
// datafusion/physical-plan/src/windows/window_agg_exec.rs:255-263
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
    if self.partition_keys().is_empty() {
        InputDistributionRequirements::new(vec![Distribution::SinglePartition])
    } else {
        InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
            self.partition_keys(),
        )])
    }
}

A global sort wants one partition:

rust
// datafusion/physical-plan/src/sorts/sort.rs:1488-1495
vec![Distribution::SinglePartition]

And a SortPreservingMergeExec explicitly wants nothing, because it consumes multiple partitions and merges them:

rust
// datafusion/physical-plan/src/sorts/sort_preserving_merge.rs:248-252
fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
    crate::InputDistributionRequirements::new(vec![
        Distribution::UnspecifiedDistribution,
    ])
}

Note the discipline here: an operator only declares what it needs. Distribution::UnspecifiedDistribution means “I don’t care” — and that is the hook the optimizer uses to decide it’s allowed to add parallelism.

Satisfying the requirement

The work happens in one physical optimizer rule, EnsureRequirements:

rust
// datafusion/physical-optimizer/src/ensure_requirements/mod.rs:166-288
/// Optimizer rule that enforces both distribution and sorting requirements.
///
/// This rule inserts the repartitioning and sorting operators needed to
/// satisfy the requirements that operators declare via
/// [`ExecutionPlan::input_distribution_requirements`] and
/// [`ExecutionPlan::required_input_ordering`].
pub struct EnsureRequirements {}

Its pipeline is a carefully ordered set of phases:

Phase 0: normalize InterleaveExec back to UnionExec (interleaves are artifacts of Phase 2 and would break if a child later loses its partitioning).
Phase 1: join key reordering, top-down, so that a parent join’s key order matches its children’s.
Phase 2a: distribution enforcement, bottom-up.
Phase 2b: sorting enforcement, bottom-up, on the distribution-fixed plan.
Phase 3: parallelize sorts, restore order-preserving variants, push sorts down, and convert some sorts to partial sorts.

Combining distribution and sorting into one pass is not cosmetic. The module docs spell out why:

text
// datafusion/physical-optimizer/src/ensure_requirements/mod.rs:122-138
/// # Behavior: joint distribution + sorting
///
/// Resolving distribution and sorting together lets Phase 2 produce a
/// parallel sort plan in cases where the two-rule pipeline historically
/// risked a serial one. Given `Sort(DESC) ← Coalesce ← MultiPartitionSource`,
/// `EnsureRequirements` rewrites it into:
///
/// ```text
/// SortPreservingMergeExec: [a DESC]            (cheap k-way merge of sorted streams)
///   SortExec: [a DESC], preserve_partitioning=true   (N sorts run in parallel)
///     MultiPartitionSource
/// ```
///
/// Each input partition is sorted in parallel, then a `SortPreservingMergeExec`
/// at the top performs a cheap merge of pre-sorted streams. For TopK queries
/// (`fetch=K`), each parallel sort only keeps K rows per partition, so total
/// memory is `N × K` rather than coalescing the entire stream first.

Phase 2a: the actual insertion

The bottom-up pass computes, for each child, whether the requirement is already met:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:1717-1748
match &requirement {
    Distribution::SinglePartition => {
        child = add_merge_on_top(child);
    }
    Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) => {
        let child_partitions =
            child.plan.output_partitioning().partition_count();
        let partitioning_satisfied = input_distributions
            .child_satisfaction(
                child_idx,
                child.plan.as_ref(),
                ChildSatisfactionOptions::new()
                    .with_allow_subset(allow_subset_satisfy_partitioning),
            )?
            .is_satisfied();
        // ...
        let needs_repartition = if allow_subset_satisfy_partitioning {
            !partitioning_satisfied
        } else {
            !partitioning_satisfied
                || (target_partitions > child_partitions
                    && !preserve_satisfying_file_partitioning)
        };
        let should_add_repartition = hash_necessary && needs_repartition;

And when it decides a shuffle is needed, it builds one:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:1752-1801
if should_add_repartition {
    let partitioning = match child.plan.output_partitioning() {
        Partitioning::Range(range) if partitioning_satisfied => {
            match range.scale(target_partitions) {
                Some(range) => { /* keep the range layout, just widen it */ }
                None => requirement
                    .clone()
                    .create_partitioning(target_partitions),
            }
        }
        _ => {
            requirement.clone().create_partitioning(target_partitions)
        }
    };
    // When there is an existing ordering, we preserve ordering during
    // repartition. ...
    let repartition = RepartitionExec::try_new(
        Arc::clone(&child.plan),
        partitioning,
    )?
    .with_preserve_order();
    let plan = Arc::new(repartition) as _;
    child = DistributionContext::new(plan, true, vec![child]);
}

Distribution::create_partitioning is the bridge from requirement to concrete operator (partitioning.rs:952-965):

rust
pub fn create_partitioning(self, partition_count: usize) -> Partitioning {
    match self {
        Distribution::UnspecifiedDistribution => {
            Partitioning::UnknownPartitioning(partition_count)
        }
        Distribution::SinglePartition => Partitioning::UnknownPartitioning(1),
        Distribution::HashPartitioned(expr) | Distribution::KeyPartitioned(expr) => {
            Partitioning::Hash(expr, partition_count)
        }
    }
}

So a KeyPartitioned([k]) requirement becomes RepartitionExec(Hash([k], target_partitions)) — exactly the operator from the previous article.

The `hash_necessary` alignment trick

A naive implementation checks each child independently: “does the right side need a hash shuffle? Yes. Does the left? No. Then shuffle the right.” But that produces garbage — a partitioned join where one side is hash-partitioned and the other is not will silently drop rows that should have matched. The optimizer handles this by aligning the flags:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:1173-1203
let is_hash = matches!(
    requirement,
    Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_)
);
// Hash re-partitioning is necessary when the input has more than one
// partitions:
let multi_partitions = child.output_partitioning().partition_count() > 1;
let roundrobin_sensible = roundrobin_beneficial && roundrobin_beneficial_stats;
needs_alignment |= is_hash && (multi_partitions || roundrobin_sensible);
repartition_status_flags.push((
    is_hash,
    RepartitionRequirementStatus {
        requirement,
        roundrobin_beneficial,
        roundrobin_beneficial_stats,
        hash_necessary: is_hash && multi_partitions,
    },
));
// ...
// Align hash necessary flags for hash partitions to generate consistent
// hash partitions at each children:
if needs_alignment {
    for (is_hash, status) in &mut repartition_status_flags {
        if *is_hash {
            status.hash_necessary = true;
        }
    }
}

If any child of a partitioned join needs a hash shuffle, all hash-keyed children get one. Consistency beats minimality here.

Co-partitioning: making the two sides agree

Even with both sides hash-partitioned, the number of partitions and the layout must match. enforce_distribution_relationships handles the cross-child constraint:

rust
// datafusion/physical-plan/src/distribution_requirements.rs:366-393
fn compatible_co_partitioning_layout(
    first_partitioning: &Partitioning,
    other_partitioning: &Partitioning,
) -> bool {
    if first_partitioning.partition_count() == 1
        && other_partitioning.partition_count() == 1
    {
        return true;
    }
    if first_partitioning.partition_count() != other_partitioning.partition_count() {
        return false;
    }
    match (first_partitioning, other_partitioning) {
        (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true,
        (Partitioning::Range(left), Partitioning::Range(right)) => {
            left.split_points() == right.split_points()
                && left.ordering().len() == right.ordering().len()
                && left
                    .ordering()
                    .iter()
                    .zip(right.ordering())
                    .all(|(left, right)| left.options == right.options)
        }
        _ => false,
    }
}

Two Hash partitionings with the same count are considered compatible because the hash function is fixed (REPARTITION_RANDOM_STATE with seed 0, see the previous article) — so hash(k) % 8 is the same function on both sides. Two Range partitionings must have identical split points, since range layouts are data-dependent and two different samplings would put the same key in different partitions.

Round robin for pure parallelism

When the requirement is UnspecifiedDistribution — the operator doesn’t care — the optimizer may still add a round-robin repartition just to get more parallelism:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:1803-1809
Distribution::UnspecifiedDistribution => {
    if add_roundrobin {
        // Add round-robin repartitioning on top of the operator
        // to increase parallelism.
        child = add_roundrobin_on_top(child, target_partitions)?;
    }
}

But it’s gated on statistics and on the operator’s own opinion:

rust
// datafusion/physical-plan/src/execution_plan.rs:245-262
fn benefits_from_input_partitioning(&self) -> Vec<bool> {
    // By default try to maximize parallelism with more CPUs if
    // possible
    self.input_distribution_requirements()
        .per_child_distributions()
        .map(|dist| !matches!(dist, Distribution::SinglePartition))
        .collect()
}

SortExec overrides this to vec![false] — a sort is not helped by more input partitions (it needs them coalesced, not split). And the round-robin is additionally gated on the child having more rows than one batch (roundrobin_beneficial_stats, enforce_distribution.rs:1165-1172), so shuffling a 3-row table to 32 cores is avoided.

Subset satisfaction

The strictest-match policy is documented at the top of the rule:

text
// datafusion/physical-optimizer/src/ensure_requirements/mod.rs:140-147
/// # Behavior: strictest distribution match for joins
///
/// Distribution requirements are met in the strictest way. For example, a
/// hash join with keys `(a, b, c)` requires `Distribution(a, b, c)`. This
/// can in principle be satisfied by partitioning on any superset of any
/// subset of `(a, b, c)`, but this rule always partitions on the exact key
/// tuple `(a, b, c)`. This is sometimes more aggressive than strictly
/// necessary, but the strictest match helps avoid data skew in joins.

Subset satisfaction is opt-in per child and is disabled in exactly the cases where it would be unsafe:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:1648-1683
// Allow subset satisfaction when:
// 1. Current partition count >= threshold
// 2. Not a partitioned join since must use exact hash matching for joins
// 3. Not a grouping set aggregate (requires exact hash including __grouping_id)
let requires_grouping_id = matches!(&requirement,
    Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs)
    if exprs.iter().any(|expr| {
        (expr.as_ref() as &dyn Any)
            .downcast_ref::<Column>()
            .is_some_and(|col| col.name() == Aggregate::INTERNAL_GROUPING_ID)
    })
);

let allow_subset_satisfy_partitioning = (current_partitions
    >= subset_satisfaction_threshold
    || (preserve_file_partition_threshold_met
        && current_partitions < target_partitions))
    && !is_partitioned_join
    && !requires_grouping_id;

__grouping_id is the hidden column that distinguishes grouping sets in ROLLUP/CUBE. If you hash on (a) but the aggregate needs (a, __grouping_id), the partial aggregates from different grouping sets would be merged into the same hash bucket and produce wrong results. Hence the explicit exclusion.

Trade-offs and failure modes

The requirement is a declaration, not a guarantee. The optimizer satisfies it at plan time using the Partitioning metadata. If a node lies about its output partitioning, everything downstream silently produces wrong results. This is why the invariants are checked: InputDistributionRequirements::check_invariants (distribution_requirements.rs:216-240) and ExecutionPlan::check_invariants verify co-partitioning at InvariantLevel::Executable.

SinglePartition is a parallelism black hole. Any operator that requires Distribution::SinglePartition — a global sort, a global aggregate, a window function with no PARTITION BY — collapses the whole pipeline to one thread below it. The optimizer’s only defence is to parallelize the work above the collapse point: sort each partition in parallel and merge with SortPreservingMergeExec. That’s the “joint distribution + sorting” behaviour quoted earlier, and it’s the single most important performance trick in the whole rule.

Skew survives the requirement. KeyPartitioned([k]) says “same key, same partition”. It says nothing about balance. If k is skewed, one partition gets 90% of the rows and the join’s critical path is that one thread. The strictest-match policy mitigates this by hashing on the full key tuple, but a genuinely hot key is unfixable at this layer — it needs range partitioning, salting, or a different algorithm.

Co-partitioning can require a second shuffle. If the left side is already Hash([a, b]) and the right side is Hash([a]), satisfying the join’s co-partitioning requirement means re-shuffling the right side onto (a, b) — even though the join’s own key requirement ((a, b)) was already met by the left. enforce_distribution_relationships (enforce_distribution.rs:1273-1330) picks a reference child that already satisfies the layout and rewrites the others to match, ranking candidates by estimated data size so the biggest child is the one you don’t want to move.

UnspecifiedDistribution is a licence to be greedy. It’s the signal “this operator would like more parallelism”. If an operator declares it but doesn’t actually benefit, the optimizer will happily add a round-robin shuffle that costs a full data copy for nothing. That’s why benefits_from_input_partitioning exists as a separate knob, and why operators like SortExec explicitly opt out.

Bottom line

The declarative-distribution design is what makes DataFusion’s optimizer composable. New operators don’t need new planner rules; they need one method:

rust
fn input_distribution_requirements(&self) -> InputDistributionRequirements

The shuffle that follows is not an implementation detail of the operator that needed it — it’s a generic, reusable artifact (RepartitionExec) that the optimizer inserts, counts, and reuses across the plan. The cost of this abstraction is that the planner now has to reason about global consistency (aligned hash necessity, co-partitioning, subset policies) that individual operators never had to think about. The benefit is that a HashJoinExec doesn’t care whether its input came from a Parquet scan, a UNION, or a subquery — it just says “I need k co-located”, and the optimizer figures out the rest.