1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
//! Cooperative yielding for input-amplifying operators (#217).
//!
//! DataFusion's `EnsureCooperative` instruments LEAF streams only: budget
//! is consumed per batch a leaf produces. An operator that amplifies its
//! input — a cross or nested-loop join whose output is orders of magnitude
//! larger than its input, or an unnest — drains its tiny budget-aware
//! inputs in microseconds and then computes budget-free: a 5-way cross
//! join over five 100-row VALUES tables feeds an aggregate 10^10 rows
//! while consuming 5 units of budget, so its poll never yields and no
//! timeout, cancel watcher, or select! arm can ever run (measured: a 2 s
//! `tokio::time::timeout` armed around it did not fire in 7+ minutes).
//!
//! The fix is one wrapper: put a [`CooperativeExec`] on top of each
//! amplifier so budget is also consumed per OUTPUT batch. The stream then
//! returns `Pending` every ~128 batches (~1M rows), which is what makes
//! the executor's cancel watcher and every timeout real for this operator
//! class. `datafusion-proto` round-trips `CooperativeExec`, so distributed
//! fragment encoding is unaffected.
use std::sync::Arc;
use datafusion::common::Result;
use datafusion::common::config::ConfigOptions;
use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::coop::CooperativeExec;
use datafusion::physical_plan::joins::{CrossJoinExec, NestedLoopJoinExec};
use datafusion::physical_plan::unnest::UnnestExec;
/// Wraps input-amplifying operators in [`CooperativeExec`] so their output
/// participates in cooperative scheduling. See the module docs for why the
/// default leaf-only instrumentation is not enough.
#[derive(Debug, Default)]
pub struct CooperativeAmplifiers {}
impl CooperativeAmplifiers {
pub fn new() -> Self {
Self {}
}
}
fn is_amplifier(plan: &dyn ExecutionPlan) -> bool {
// `ExecutionPlan: Any` — upcast to downcast (DF 54 has no `as_any`).
let any = plan as &dyn std::any::Any;
any.downcast_ref::<CrossJoinExec>().is_some()
|| any.downcast_ref::<NestedLoopJoinExec>().is_some()
// The module docs have always named unnest as a member of this class
// and it was never actually matched: one row in, one row per list
// element out, with no leaf of its own between it and the consumer.
|| any.downcast_ref::<UnnestExec>().is_some()
}
/// Is this node already a [`CooperativeExec`]?
fn is_cooperative(plan: &dyn ExecutionPlan) -> bool {
(plan as &dyn std::any::Any)
.downcast_ref::<CooperativeExec>()
.is_some()
}
impl PhysicalOptimizerRule for CooperativeAmplifiers {
fn optimize(
&self,
plan: Arc<dyn ExecutionPlan>,
_config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
plan.transform_up(|node| {
if is_amplifier(node.as_ref()) {
return Ok(Transformed::yes(
Arc::new(CooperativeExec::new(node)) as Arc<dyn ExecutionPlan>
));
}
// Collapse a doubled wrapper, so applying the rule to a plan it has
// already run on is a no-op. `transform_up` visits children first:
// an amplifier that was *already* wrapped gets a second wrapper
// when we reach it, and the pre-existing one is then visited with
// that as its child. Without this the plan would gain a layer per
// pass, and each layer costs a poll indirection on every batch.
if is_cooperative(node.as_ref())
&& node
.children()
.first()
.is_some_and(|child| is_cooperative(child.as_ref()))
&& let Some(child) = node.children().first()
{
return Ok(Transformed::yes(Arc::clone(child)));
}
Ok(Transformed::no(node))
})
.map(|t| t.data)
}
fn name(&self) -> &str {
"CooperativeAmplifiers"
}
fn schema_check(&self) -> bool {
// A CooperativeExec wrapper is schema-transparent.
true
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::physical_plan::displayable;
fn leaf() -> Arc<dyn ExecutionPlan> {
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
MemorySourceConfig::try_new_exec(&[vec![]], schema, None).unwrap()
}
fn optimize(plan: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
CooperativeAmplifiers::new()
.optimize(plan, &ConfigOptions::default())
.unwrap()
}
fn rendered(plan: &Arc<dyn ExecutionPlan>) -> String {
format!("{}", displayable(plan.as_ref()).indent(false))
}
/// The operator class the rule exists for must actually be wrapped.
#[test]
fn a_cross_join_is_wrapped() {
let join = Arc::new(CrossJoinExec::new(leaf(), leaf())) as Arc<dyn ExecutionPlan>;
assert_eq!(rendered(&join).matches("Cooperative").count(), 0);
let out = optimize(join);
assert_eq!(
rendered(&out).matches("Cooperative").count(),
1,
"cross join not wrapped:\n{}",
rendered(&out)
);
}
/// Running the rule on a plan it has already run on must change nothing.
///
/// Without the collapse, `transform_up` adds a wrapper every pass: it
/// visits the amplifier before the wrapper already above it, so the old
/// wrapper simply ends up on top of the new one. Each layer is a poll
/// indirection on every batch for the rest of the query.
#[test]
fn optimizing_twice_is_a_no_op() {
let join = Arc::new(CrossJoinExec::new(leaf(), leaf())) as Arc<dyn ExecutionPlan>;
let once = optimize(join);
let twice = optimize(Arc::clone(&once));
assert_eq!(
rendered(&once),
rendered(&twice),
"the rule is not idempotent; a second pass added a layer"
);
assert_eq!(
rendered(&twice).matches("Cooperative").count(),
1,
"expected exactly one wrapper:\n{}",
rendered(&twice)
);
}
/// A plan with nothing to amplify must come back untouched — the rule
/// costs a poll indirection, so it should only be paid where it buys
/// preemptibility.
#[test]
fn a_plan_without_an_amplifier_is_left_alone() {
let plan = leaf();
let out = optimize(Arc::clone(&plan));
assert_eq!(rendered(&plan), rendered(&out));
assert_eq!(rendered(&out).matches("Cooperative").count(), 0);
}
/// Plan `sql` against two small tables and return the physical plan.
///
/// Built through the planner rather than by calling operator constructors
/// directly: the point is that the shape DataFusion actually emits for
/// these queries is matched, which a hand-built node cannot tell us.
async fn physical(sql: &str) -> Arc<dyn ExecutionPlan> {
use datafusion::arrow::array::Int64Array;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::datasource::MemTable;
use datafusion::prelude::SessionContext;
let ctx = SessionContext::new();
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
for name in ["t1", "t2"] {
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(Int64Array::from(vec![1i64, 2, 3]))],
)
.unwrap();
let table = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]]).unwrap();
ctx.register_table(name, Arc::new(table)).unwrap();
}
let logical = ctx.sql(sql).await.unwrap().into_optimized_plan().unwrap();
ctx.state().create_physical_plan(&logical).await.unwrap()
}
/// A nested-loop join is named in `is_amplifier` and was never tested.
///
/// It is the *other* operator the module was written for — a non-equi join
/// whose output dwarfs its input — and until now only the cross join had a
/// test, so this arm could have been deleted without anything failing.
#[tokio::test]
async fn a_nested_loop_join_is_wrapped() {
let plan = physical("SELECT t1.a FROM t1, t2 WHERE t1.a < t2.a").await;
assert!(
rendered(&plan).contains("NestedLoopJoin"),
"fixture stopped producing a nested-loop join:\n{}",
rendered(&plan)
);
let out = optimize(plan);
assert!(
rendered(&out).contains("Cooperative"),
"a nested-loop join must be made preemptible:\n{}",
rendered(&out)
);
}
/// Unnest is the arm that already regressed once.
///
/// The module docs "have always named unnest as a member of this class and
/// it was never actually matched" — a silent gap that survived because no
/// test covered it. This is that test.
#[tokio::test]
async fn an_unnest_is_wrapped() {
let plan = physical("SELECT unnest([1, 2, 3]) AS u FROM t1").await;
assert!(
rendered(&plan).contains("Unnest"),
"fixture stopped producing an unnest:\n{}",
rendered(&plan)
);
let out = optimize(plan);
assert!(
rendered(&out).contains("Cooperative"),
"an unnest must be made preemptible:\n{}",
rendered(&out)
);
}
/// Idempotence must hold for every amplifier, not just the cross join.
///
/// The collapse only inspects `children().first()`, so an operator whose
/// wrapped form sits differently in the tree would grow a layer per pass —
/// invisible until a plan had been optimized twice in production.
#[tokio::test]
async fn every_amplifier_is_idempotent_under_a_second_pass() {
for sql in [
"SELECT t1.a FROM t1, t2 WHERE t1.a < t2.a",
"SELECT unnest([1, 2, 3]) AS u FROM t1",
] {
let once = optimize(physical(sql).await);
let twice = optimize(Arc::clone(&once));
assert_eq!(
rendered(&once),
rendered(&twice),
"a second pass changed the plan for:\n{sql}"
);
}
}
}