polars_stream/
skeleton.rs

1#![allow(unused)] // TODO: remove me
2use std::cmp::Reverse;
3use std::time::{Duration, Instant};
4
5use parking_lot::Mutex;
6use polars_core::prelude::*;
7use polars_core::query_result::QueryResult;
8use polars_core::runtime::RAYON;
9use polars_expr::planner::{ExpressionConversionState, create_physical_expr, get_expr_depth_limit};
10use polars_plan::plans::{IR, IRPlan, IRPlanSorted};
11use polars_plan::prelude::AExpr;
12use polars_plan::prelude::expr_ir::ExprIR;
13use polars_utils::arena::{Arena, Node};
14use polars_utils::relaxed_cell::RelaxedCell;
15use slotmap::{SecondaryMap, SlotMap};
16
17use crate::graph::{Graph, GraphNodeKey};
18use crate::metrics::GraphMetrics;
19use crate::physical_plan::{PhysNode, PhysNodeKey, PhysNodeKind, StreamingLowerIRContext};
20
21/// Executes the IR with the streaming engine.
22///
23/// Unsupported operations can fall back to the in-memory engine.
24///
25/// Returns:
26/// - `Ok(QueryResult::Single(DataFrame))` when collecting to a single sink.
27/// - `Ok(QueryResult::Multiple(Vec<DataFrame>))` when collecting to multiple sinks.
28/// - `Err` if the IR can't be executed.
29///
30/// Returned `DataFrame`s contain data only for memory sinks,
31/// `DataFrame`s corresponding to file sinks are empty.
32pub fn run_query(
33    node: Node,
34    ir_arena: &mut Arena<IR>,
35    expr_arena: &mut Arena<AExpr>,
36) -> PolarsResult<QueryResult> {
37    StreamingQuery::build(node, ir_arena, expr_arena)?.execute()
38}
39
40/// Visualizes the physical plan as a dot graph.
41pub fn visualize_physical_plan(
42    node: Node,
43    ir_arena: &mut Arena<IR>,
44    expr_arena: &mut Arena<AExpr>,
45) -> PolarsResult<String> {
46    let mut phys_sm = SlotMap::with_capacity_and_key(ir_arena.len());
47    let sortedness = IRPlanSorted::resolve(node, ir_arena, expr_arena);
48
49    let ctx = StreamingLowerIRContext {
50        prepare_visualization: true,
51        sortedness: &sortedness,
52    };
53    let root_phys_node =
54        crate::physical_plan::build_physical_plan(node, ir_arena, expr_arena, &mut phys_sm, ctx)?;
55
56    let out = crate::physical_plan::visualize_plan(root_phys_node, &phys_sm, expr_arena);
57
58    Ok(out)
59}
60
61pub struct StreamingQuery {
62    top_ir: IR,
63    pub graph: Graph,
64    pub root_phys_node: PhysNodeKey,
65    pub phys_sm: SlotMap<PhysNodeKey, PhysNode>,
66    pub phys_to_graph: SecondaryMap<PhysNodeKey, GraphNodeKey>,
67    pub metrics: Option<Arc<Mutex<GraphMetrics>>>,
68}
69
70/// Configures if IR lowering creates the `format_str` for `InMemoryMap`.
71pub static PREPARE_VISUALIZATION_DATA: RelaxedCell<bool> = RelaxedCell::new_bool(false);
72
73/// Sets config to ensure IR lowering always creates the `format_str` for `InMemoryMap`.
74pub fn always_prepare_visualization_data() {
75    PREPARE_VISUALIZATION_DATA.store(true);
76}
77
78fn cfg_prepare_visualization_data() -> bool {
79    if !PREPARE_VISUALIZATION_DATA.load() {
80        PREPARE_VISUALIZATION_DATA.fetch_or(
81            std::env::var("POLARS_STREAM_ALWAYS_PREPARE_VISUALIZATION_DATA").as_deref() == Ok("1"),
82        );
83    }
84
85    PREPARE_VISUALIZATION_DATA.load()
86}
87
88impl StreamingQuery {
89    pub fn build(
90        node: Node,
91        ir_arena: &mut Arena<IR>,
92        expr_arena: &mut Arena<AExpr>,
93    ) -> PolarsResult<Self> {
94        if let Ok(visual_path) = std::env::var("POLARS_VISUALIZE_IR") {
95            let plan = IRPlan {
96                lp_top: node,
97                lp_arena: ir_arena.clone(),
98                expr_arena: expr_arena.clone(),
99            };
100            let visualization = plan.display_dot().to_string();
101            std::fs::write(visual_path, visualization).unwrap();
102        }
103        let mut phys_sm = SlotMap::with_capacity_and_key(ir_arena.len());
104        let sortedness = IRPlanSorted::resolve(node, ir_arena, expr_arena);
105        let ctx = StreamingLowerIRContext {
106            prepare_visualization: cfg_prepare_visualization_data(),
107            sortedness: &sortedness,
108        };
109        let root_phys_node = crate::physical_plan::build_physical_plan(
110            node,
111            ir_arena,
112            expr_arena,
113            &mut phys_sm,
114            ctx,
115        )?;
116        if let Ok(visual_path) = std::env::var("POLARS_VISUALIZE_PHYSICAL_PLAN") {
117            let visualization =
118                crate::physical_plan::visualize_plan(root_phys_node, &phys_sm, expr_arena);
119            std::fs::write(visual_path, visualization).unwrap();
120        }
121
122        let (mut graph, phys_to_graph) =
123            crate::physical_plan::physical_plan_to_graph(root_phys_node, &phys_sm, expr_arena)?;
124
125        let top_ir = ir_arena.get(node).clone();
126
127        let metrics = if std::env::var("POLARS_TRACK_METRICS").as_deref() == Ok("1")
128            || std::env::var("POLARS_LOG_METRICS").as_deref() == Ok("1")
129        {
130            polars_async::executor::track_task_metrics(true);
131            Some(Arc::default())
132        } else {
133            None
134        };
135
136        let out = StreamingQuery {
137            top_ir,
138            graph,
139            root_phys_node,
140            phys_sm,
141            phys_to_graph,
142            metrics,
143        };
144
145        Ok(out)
146    }
147
148    pub fn execute(self) -> PolarsResult<QueryResult> {
149        let StreamingQuery {
150            top_ir,
151            mut graph,
152            root_phys_node,
153            phys_sm,
154            phys_to_graph,
155            metrics,
156        } = self;
157
158        let query_start = Instant::now();
159        let mut results = crate::execute::execute_graph(&mut graph, metrics.clone())?;
160        let query_elapsed = query_start.elapsed();
161
162        // Print metrics.
163        if let Some(lock) = metrics
164            && std::env::var("POLARS_LOG_METRICS").as_deref() == Ok("1")
165        {
166            let mut total_query_ns = 0;
167            let mut lines = Vec::new();
168            let m = lock.lock();
169            for phys_node_key in phys_sm.keys() {
170                let Some(graph_node_key) = phys_to_graph.get(phys_node_key) else {
171                    continue;
172                };
173                let Some(node_metrics) = m.get(*graph_node_key) else {
174                    continue;
175                };
176                let name = graph.nodes[*graph_node_key].compute.name();
177                let total_ns =
178                    node_metrics.total_poll_time_ns + node_metrics.total_state_update_time_ns;
179                let total_time = Duration::from_nanos(total_ns);
180                let poll_time = Duration::from_nanos(node_metrics.total_poll_time_ns);
181                let update_time = Duration::from_nanos(node_metrics.total_state_update_time_ns);
182                let max_poll_time = Duration::from_nanos(node_metrics.max_poll_time_ns);
183                let max_update_time = Duration::from_nanos(node_metrics.max_state_update_time_ns);
184                let total_polls = node_metrics.total_polls;
185                let total_updates = node_metrics.total_state_updates;
186                let perc_stolen = node_metrics.total_stolen_polls as f64
187                    / node_metrics.total_polls as f64
188                    * 100.0;
189
190                let rows_received = node_metrics.rows_received;
191                let morsels_received = node_metrics.morsels_received;
192                let max_received = node_metrics.largest_morsel_received;
193                let rows_sent = node_metrics.rows_sent;
194                let morsels_sent = node_metrics.morsels_sent;
195                let max_sent = node_metrics.largest_morsel_sent;
196
197                let io_total_active_time = Duration::from_nanos(node_metrics.io_total_active_ns);
198                let io_total_bytes_requested = node_metrics.io_total_bytes_requested;
199                let io_total_bytes_received = node_metrics.io_total_bytes_received;
200                let io_total_bytes_sent = node_metrics.io_total_bytes_sent;
201
202                lines.push(
203                    (total_time, format!(
204                        "{name}: tot({total_time:.2?}), \
205                                 poll({poll_time:.2?}, n={total_polls}, max={max_poll_time:.2?}, stolen={perc_stolen:.1}%), \
206                                 update({update_time:.2?}, n={total_updates}, max={max_update_time:.2?}), \
207                                 recv(row={rows_received}, morsel={morsels_received}, max={max_received}), \
208                                 sent(row={rows_sent}, morsel={morsels_sent}, max={max_sent}), \
209                                 io(\
210                                    total_active_time={io_total_active_time:.2?}, \
211                                    total_bytes_requested={io_total_bytes_requested}, \
212                                    total_bytes_received={io_total_bytes_received}, \
213                                    total_bytes_sent={io_total_bytes_sent})"))
214                );
215
216                total_query_ns += total_ns;
217            }
218            lines.sort_by_key(|(tot, _)| Reverse(*tot));
219
220            let total_query_time = Duration::from_nanos(total_query_ns);
221            eprintln!(
222                "Streaming query took {query_elapsed:.2?} ({total_query_time:.2?} CPU), detailed breakdown:"
223            );
224            for (_tot, line) in lines {
225                eprintln!("{line}");
226            }
227            eprintln!();
228        }
229
230        match top_ir {
231            IR::SinkMultiple { inputs } => {
232                let phys_node = &phys_sm[root_phys_node];
233                let PhysNodeKind::SinkMultiple { sinks } = phys_node.kind() else {
234                    unreachable!();
235                };
236
237                Ok(QueryResult::Multiple(
238                    sinks
239                        .iter()
240                        .map(|phys_node_key| {
241                            results
242                                .remove(phys_to_graph[*phys_node_key])
243                                .unwrap_or_else(DataFrame::empty)
244                        })
245                        .collect(),
246                ))
247            },
248            _ => Ok(QueryResult::Single(
249                results
250                    .remove(phys_to_graph[root_phys_node])
251                    .unwrap_or_else(DataFrame::empty),
252            )),
253        }
254    }
255}
polars_stream/skeleton.rs

polars_stream/
skeleton.rs
