krishiv_plan/optimizer/auto_partition.rs
1//! AQE auto-partition rule.
2
3use crate::{Partitioning, PhysicalPlan};
4
5use super::{AqeRule, RuntimeStats, StreamingAqeGuard};
6
7const DEFAULT_TARGET_PARTITION_BYTES: u64 = krishiv_common::partition::TARGET_BYTES_PER_PARTITION;
8
9/// AQE rule that adjusts the bucket count of `Hash` and `RoundRobin` exchange
10/// nodes based on the observed data volume from the previous execution.
11///
12/// The rule reads `RuntimeStats` (one per DataFusion partition), sums
13/// `memory_bytes` to obtain the total stage output size, and computes a target
14/// partition count:
15///
16/// `target = clamp(1, max_buckets, ceil(total_bytes / target_partition_bytes))`
17///
18/// The target is applied unconditionally: the rule can both increase and
19/// decrease bucket counts. This matches Spark AQE's behavior — if early
20/// execution stages produced far less data than expected, the rule shrinks
21/// the downstream partition count to avoid over-parallelism (task scheduling
22/// overhead dominating actual work). The minimum floor is always 1.
23///
24/// When stats are empty (first execution) or contain no measurable memory, the
25/// rule is a no-op and returns `None`.
26pub struct AutoPartitionRule {
27 /// Desired bytes per partition. Default: 128 MiB.
28 target_partition_bytes: u64,
29 /// Upper bound on the number of partitions. Derived from
30 /// `target_partitions` in the session config so we never ask for more
31 /// parallelism than the runtime can supply.
32 max_buckets: u32,
33}
34
35impl AutoPartitionRule {
36 /// Create a new rule with the given max bucket count.
37 ///
38 /// Uses the default `target_partition_bytes` of 128 MiB.
39 pub fn new(max_buckets: u32) -> Self {
40 Self {
41 target_partition_bytes: DEFAULT_TARGET_PARTITION_BYTES,
42 max_buckets,
43 }
44 }
45
46 /// Set a custom `target_partition_bytes`.
47 #[must_use]
48 pub fn with_target_partition_bytes(mut self, bytes: u64) -> Self {
49 self.target_partition_bytes = bytes;
50 self
51 }
52}
53
54impl AqeRule for AutoPartitionRule {
55 fn name(&self) -> &str {
56 "auto-partition"
57 }
58
59 fn apply(&self, plan: &PhysicalPlan, stats: &[RuntimeStats]) -> Option<PhysicalPlan> {
60 // When an explicit shuffle_partitions override is set on the plan
61 // (via SET shuffle.partitions = N or SessionBuilder), use it as the
62 // target bucket count regardless of stats. Stats may be empty on the
63 // first execution and that's fine — the override is a user intent.
64 if let Some(override_buckets) = plan.shuffle_partitions() {
65 return self.apply_override(plan, override_buckets);
66 }
67
68 if stats.is_empty() || StreamingAqeGuard::plan_is_streaming(plan) {
69 return None;
70 }
71
72 // Sum the best available size metric across all partitions.
73 // Prefer serialized_bytes (shuffle wire size) over memory_bytes (peak
74 // in-memory) because shuffle output is compressed/serialized and thus
75 // a more accurate proxy for partition cost. Fall back to memory_bytes
76 // when serialized_bytes is zero (non-shuffle tasks or older executors).
77 let total_bytes: u64 = stats
78 .iter()
79 .map(|s| {
80 if s.serialized_bytes > 0 {
81 s.serialized_bytes
82 } else {
83 s.memory_bytes
84 }
85 })
86 .sum();
87 if total_bytes == 0 {
88 return None;
89 }
90
91 // Compute target partition count via the shared cross-mode sizing brain.
92 let target = krishiv_common::partition::recommend_buckets(
93 total_bytes,
94 1,
95 self.max_buckets,
96 self.target_partition_bytes,
97 );
98
99 self.stamp_target(plan, target)
100 }
101}
102
103impl AutoPartitionRule {
104 /// Apply the rule with an explicit override bucket count.
105 /// Skips streaming plans, but does not require runtime stats.
106 fn apply_override(&self, plan: &PhysicalPlan, target: u32) -> Option<PhysicalPlan> {
107 if StreamingAqeGuard::plan_is_streaming(plan) {
108 return None;
109 }
110 let target = target.max(1);
111 self.stamp_target(plan, target)
112 }
113
114 /// Stamp `target` bucket count onto all Hash/RoundRobin exchange nodes
115 /// whose current count differs from `target`. Returns `None` if no node
116 /// needed adjustment.
117 ///
118 /// Both increases and decreases are applied — if the observed data volume
119 /// implies fewer partitions than currently planned, the bucket count is
120 /// lowered to avoid over-parallelism (task scheduling overhead > useful
121 /// work). The caller guarantees `target >= 1`.
122 fn stamp_target(&self, plan: &PhysicalPlan, target: u32) -> Option<PhysicalPlan> {
123 let mut changed = false;
124 for node in plan.nodes() {
125 match node.partitioning() {
126 Partitioning::Hash { buckets, .. } | Partitioning::RoundRobin { buckets, .. }
127 if *buckets != target =>
128 {
129 changed = true;
130 }
131 _ => {}
132 }
133 }
134
135 if !changed {
136 return None;
137 }
138
139 // Only clone when we know a rewrite is needed.
140 let mut plan = plan.clone();
141 for node in plan.nodes_mut() {
142 let old = node.partitioning().clone();
143 match old {
144 Partitioning::Hash { ref keys, buckets } if buckets != target => {
145 node.set_partitioning(Partitioning::Hash {
146 keys: keys.clone(),
147 buckets: target,
148 });
149 }
150 Partitioning::RoundRobin { buckets } if buckets != target => {
151 node.set_partitioning(Partitioning::RoundRobin { buckets: target });
152 }
153 _ => {}
154 }
155 }
156
157 tracing::debug!(rule = "auto-partition", target, "AutoPartitionRule applied");
158
159 Some(plan)
160 }
161}