krishiv_sql/
coop_amplifiers.rs

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

krishiv_sql/
coop_amplifiers.rs
