Struct AutoPartitionRule 

pub struct AutoPartitionRule { /* private fields */ }

AQE rule that adjusts the bucket count of Hash and RoundRobin exchange nodes based on the observed data volume from the previous execution.

The rule reads RuntimeStats (one per DataFusion partition), sums memory_bytes to obtain the total stage output size, and computes a target partition count:

target = clamp(1, max_buckets, ceil(total_bytes / target_partition_bytes))

The target is applied unconditionally: the rule can both increase and decrease bucket counts. This matches Spark AQE’s behavior — if early execution stages produced far less data than expected, the rule shrinks the downstream partition count to avoid over-parallelism (task scheduling overhead dominating actual work). The minimum floor is always 1.

When stats are empty (first execution) or contain no measurable memory, the rule is a no-op and returns None.

Implementations

impl AutoPartitionRule

pub fn new(max_buckets: u32) -> Self

Create a new rule with the given max bucket count.

Uses the default target_partition_bytes of 128 MiB.

pub fn with_target_partition_bytes(self, bytes: u64) -> Self

Set a custom target_partition_bytes.

Trait Implementations

impl AqeRule for AutoPartitionRule

fn name(&self) -> &str

Short, stable rule name used in explain and diagnostics output.

fn apply( &self, plan: &PhysicalPlan, stats: &[RuntimeStats], ) -> Option<PhysicalPlan>

Apply the AQE rule given collected RuntimeStats for each stage.