ursa_engine/
serial.rs

1//! Plans that run on one task, for the browser (#278).
2//!
3//! DataFusion runs each input partition of a [`CoalescePartitionsExec`] as its
4//! own Tokio task, and so does `execute_stream` over a plan with several
5//! partitions. The browser build has no Tokio runtime: its futures are driven
6//! by JavaScript's event loop, so a spawn panics ("there is no reactor
7//! running"). `target_partitions = 1` keeps scans and joins to one partition,
8//! but a `UNION` still has one partition per input, and a graph's `nodes`
9//! table is a union.
10//!
11//! [`SingleTask`] is a physical optimizer rule, run last in the browser build.
12//! It replaces each coalesce over several partitions, and a plan whose output
13//! has several partitions, with [`SerialCoalesceExec`], which reads the
14//! partitions one after another on the task that polls it. A coalesce returns
15//! rows in no fixed order (whichever partition is ready first), so reading
16//! them in partition order is one of the orders it may produce. Graph inputs
17//! from such plans are sorted by `KeepEdgeOrder` before the index is built, so
18//! graph results do not depend on it either.
19
20use std::fmt;
21use std::sync::Arc;
22
23use datafusion::common::config::ConfigOptions;
24use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
25use datafusion::error::{DataFusionError, Result};
26use datafusion::execution::TaskContext;
27use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
28use datafusion::physical_optimizer::PhysicalOptimizerRule;
29use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
30use datafusion::physical_plan::repartition::RepartitionExec;
31use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
32use datafusion::physical_plan::{
33    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning,
34    PlanProperties, SendableRecordBatchStream,
35};
36use futures::StreamExt;
37
38/// The rule that makes every plan run on one task. See the module docs.
39#[derive(Debug, Default)]
40pub struct SingleTask;
41
42impl PhysicalOptimizerRule for SingleTask {
43    fn optimize(
44        &self,
45        plan: Arc<dyn ExecutionPlan>,
46        _config: &ConfigOptions,
47    ) -> Result<Arc<dyn ExecutionPlan>> {
48        let plan = plan
49            .transform_up(|node| {
50                if node.downcast_ref::<RepartitionExec>().is_some() {
51                    // Not planned with one target partition; it would spawn.
52                    return Err(DataFusionError::Internal(format!(
53                        "the browser build cannot run a repartition: {}",
54                        node.name()
55                    )));
56                }
57                match node.downcast_ref::<CoalescePartitionsExec>() {
58                    Some(coalesce) if partitions(coalesce.input()) > 1 => Ok(Transformed::yes(
59                        Arc::new(SerialCoalesceExec::new(
60                            Arc::clone(coalesce.input()),
61                            coalesce.fetch(),
62                        )) as Arc<dyn ExecutionPlan>,
63                    )),
64                    _ => Ok(Transformed::no(node)),
65                }
66            })?
67            .data;
68        if partitions(&plan) > 1 {
69            return Ok(Arc::new(SerialCoalesceExec::new(plan, None)));
70        }
71        Ok(plan)
72    }
73
74    fn name(&self) -> &str {
75        "single_task"
76    }
77
78    fn schema_check(&self) -> bool {
79        true
80    }
81}
82
83fn partitions(plan: &Arc<dyn ExecutionPlan>) -> usize {
84    plan.output_partitioning().partition_count()
85}
86
87/// One partition holding every row of its input's partitions, read in
88/// partition order on the polling task, with an optional row limit.
89#[derive(Debug)]
90pub struct SerialCoalesceExec {
91    input: Arc<dyn ExecutionPlan>,
92    fetch: Option<usize>,
93    properties: Arc<PlanProperties>,
94}
95
96impl SerialCoalesceExec {
97    pub fn new(input: Arc<dyn ExecutionPlan>, fetch: Option<usize>) -> Self {
98        let properties = Arc::new(PlanProperties::new(
99            EquivalenceProperties::new(input.schema()),
100            Partitioning::UnknownPartitioning(1),
101            input.properties().emission_type,
102            input.properties().boundedness,
103        ));
104        Self {
105            input,
106            fetch,
107            properties,
108        }
109    }
110}
111
112impl DisplayAs for SerialCoalesceExec {
113    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
114        write!(f, "SerialCoalesceExec")?;
115        if let Some(fetch) = self.fetch {
116            write!(f, ": fetch={fetch}")?;
117        }
118        Ok(())
119    }
120}
121
122impl ExecutionPlan for SerialCoalesceExec {
123    fn name(&self) -> &str {
124        "SerialCoalesceExec"
125    }
126
127    fn properties(&self) -> &Arc<PlanProperties> {
128        &self.properties
129    }
130
131    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
132        vec![&self.input]
133    }
134
135    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
136        vec![false]
137    }
138
139    fn apply_expressions(
140        &self,
141        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
142    ) -> Result<TreeNodeRecursion> {
143        Ok(TreeNodeRecursion::Continue)
144    }
145
146    fn with_new_children(
147        self: Arc<Self>,
148        mut children: Vec<Arc<dyn ExecutionPlan>>,
149    ) -> Result<Arc<dyn ExecutionPlan>> {
150        match (children.pop(), children.is_empty()) {
151            (Some(input), true) => Ok(Arc::new(Self::new(input, self.fetch))),
152            _ => Err(DataFusionError::Internal(
153                "SerialCoalesceExec has one child".into(),
154            )),
155        }
156    }
157
158    fn execute(
159        &self,
160        partition: usize,
161        context: Arc<TaskContext>,
162    ) -> Result<SendableRecordBatchStream> {
163        if partition != 0 {
164            return Err(DataFusionError::Internal(format!(
165                "SerialCoalesceExec has one partition; asked for {partition}"
166            )));
167        }
168        let input = Arc::clone(&self.input);
169        // Each partition starts only when the one before it is finished.
170        let batches = futures::stream::iter(0..partitions(&input))
171            .flat_map(move |p| match input.execute(p, Arc::clone(&context)) {
172                Ok(stream) => stream.left_stream(),
173                Err(e) => futures::stream::once(async move { Err(e) }).right_stream(),
174            })
175            .scan(self.fetch, |remaining, batch| {
176                let item = match (batch, remaining.as_mut()) {
177                    (Err(e), _) => Some(Err(e)),
178                    (Ok(batch), None) => Some(Ok(batch)),
179                    (Ok(_), Some(0)) => None,
180                    (Ok(batch), Some(left)) => {
181                        let take = (*left).min(batch.num_rows());
182                        *left -= take;
183                        Some(Ok(batch.slice(0, take)))
184                    }
185                };
186                futures::future::ready(item)
187            });
188        Ok(Box::pin(RecordBatchStreamAdapter::new(
189            self.input.schema(),
190            batches,
191        )))
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use datafusion::arrow::util::pretty::pretty_format_batches;
199    use datafusion::physical_plan::collect;
200    use datafusion::prelude::{SessionConfig, SessionContext};
201
202    /// Run `sql` with several target partitions, with and without the rule,
203    /// and return both results printed.
204    async fn both(sql: &str) -> (String, String, Arc<dyn ExecutionPlan>) {
205        let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(4));
206        let plan = ctx
207            .sql(sql)
208            .await
209            .unwrap()
210            .create_physical_plan()
211            .await
212            .unwrap();
213        let expected =
214            pretty_format_batches(&collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap())
215                .unwrap()
216                .to_string();
217        let serial = SingleTask
218            .optimize(plan, &ConfigOptions::default())
219            .unwrap();
220        let got =
221            pretty_format_batches(&collect(Arc::clone(&serial), ctx.task_ctx()).await.unwrap())
222                .unwrap()
223                .to_string();
224        (expected, got, serial)
225    }
226
227    fn has_coalesce_over_many(plan: &Arc<dyn ExecutionPlan>) -> bool {
228        plan.exists(|node| {
229            Ok(node
230                .downcast_ref::<CoalescePartitionsExec>()
231                .is_some_and(|c| partitions(c.input()) > 1))
232        })
233        .unwrap()
234    }
235
236    #[tokio::test]
237    async fn a_union_runs_serially_with_the_same_rows() {
238        // Sorted, so the parallel and serial runs must agree exactly.
239        let (expected, got, plan) =
240            both("SELECT x FROM (SELECT 1 AS x UNION ALL SELECT 2 UNION ALL SELECT 3) ORDER BY x")
241                .await;
242        assert_eq!(expected, got);
243        assert_eq!(partitions(&plan), 1);
244        assert!(!has_coalesce_over_many(&plan));
245    }
246
247    #[tokio::test]
248    async fn a_limit_over_a_union_keeps_its_fetch() {
249        let (_, got, plan) =
250            both("SELECT * FROM (SELECT 1 AS x UNION ALL SELECT 2 UNION ALL SELECT 3) LIMIT 2")
251                .await;
252        assert_eq!(got.lines().filter(|l| l.starts_with("| ")).count(), 1 + 2);
253        assert_eq!(partitions(&plan), 1);
254    }
255
256    #[tokio::test]
257    async fn a_multi_partition_root_gets_one_partition() {
258        let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(4));
259        let plan = ctx
260            .sql("SELECT 1 AS x UNION ALL SELECT 2")
261            .await
262            .unwrap()
263            .create_physical_plan()
264            .await
265            .unwrap();
266        assert!(partitions(&plan) > 1);
267        let serial = SingleTask
268            .optimize(plan, &ConfigOptions::default())
269            .unwrap();
270        assert_eq!(partitions(&serial), 1);
271        let rows: usize = collect(serial, ctx.task_ctx())
272            .await
273            .unwrap()
274            .iter()
275            .map(|b| b.num_rows())
276            .sum();
277        assert_eq!(rows, 2);
278    }
279}
ursa_engine/serial.rs

ursa_engine/
serial.rs
