oxidelake_runtime/
dashboard.rs

1//! Building a [`DashboardModel`] from a live embedded session: the executed
2//! plan as a DAG with placement tags, per-operator counters aggregated from
3//! the session's [`TelemetryHub`](oxidelake_core::telemetry::TelemetryHub), and
4//! per-column profiles of the registered
5//! tables. This is what `oxide tui --query …` renders; without a query the CLI
6//! falls back to the synthetic demo fixture.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use datafusion::arrow::array::RecordBatch;
12use datafusion::arrow::datatypes::DataType;
13use datafusion::arrow::util::display::{ArrayFormatter, FormatOptions};
14use datafusion::physical_plan::{ExecutionPlan, collect, displayable};
15use oxidelake_compute::{GpuAggregateExec, GpuFilterExec, GpuHashJoinExec, GpuVectorDistanceExec};
16use oxidelake_core::telemetry::{OperatorSnapshot, PlanNodeSummary};
17use oxidelake_core::{BackendKind, EngineError};
18use oxidelake_tui::{ColumnProfile, DashboardModel};
19
20use crate::session::OxideSession;
21
22/// The backend a physical node was planned for: the tag on our `Gpu*Exec`
23/// nodes, CPU for every stock DataFusion node.
24fn node_backend(node: &dyn ExecutionPlan) -> BackendKind {
25    if let Some(exec) = node.downcast_ref::<GpuFilterExec>() {
26        exec.target()
27    } else if let Some(exec) = node.downcast_ref::<GpuHashJoinExec>() {
28        exec.target()
29    } else if let Some(exec) = node.downcast_ref::<GpuAggregateExec>() {
30        exec.target()
31    } else if let Some(exec) = node.downcast_ref::<GpuVectorDistanceExec>() {
32        exec.target()
33    } else {
34        BackendKind::CpuSimd
35    }
36}
37
38/// The physical plan as the dashboard's DAG rows (pre-order, with depths).
39pub fn plan_summary(plan: &Arc<dyn ExecutionPlan>) -> Vec<PlanNodeSummary> {
40    fn walk(node: &Arc<dyn ExecutionPlan>, depth: usize, out: &mut Vec<PlanNodeSummary>) {
41        let line = displayable(node.as_ref()).one_line().to_string();
42        let detail = line
43            .split_once(": ")
44            .map_or(String::new(), |(_, rest)| rest.trim().to_owned());
45        out.push(PlanNodeSummary {
46            depth,
47            name: node.name().to_owned(),
48            backend: node_backend(node.as_ref()),
49            detail,
50        });
51        for child in node.children() {
52            walk(child, depth + 1, out);
53        }
54    }
55    let mut out = Vec::new();
56    walk(plan, 0, &mut out);
57    out
58}
59
60/// Aggregates the hub's per-partition operator counters by operator name and
61/// aligns them with `plan` (the Inspector panel pairs plan row *i* with
62/// operator *i*). Stock DataFusion nodes do not report into the hub and show
63/// zeros.
64fn operators_for_plan(
65    plan: &[PlanNodeSummary],
66    recorded: &[OperatorSnapshot],
67) -> Vec<OperatorSnapshot> {
68    let mut by_name: HashMap<&str, OperatorSnapshot> = HashMap::new();
69    for op in recorded {
70        let entry = by_name
71            .entry(op.name.as_str())
72            .or_insert_with(|| OperatorSnapshot {
73                id: 0,
74                name: op.name.clone(),
75                backend: op.backend,
76                rows_in: 0,
77                rows_out: 0,
78                batches: 0,
79                elapsed_ns: 0,
80                bytes_h2d: 0,
81                bytes_d2h: 0,
82                memory_bytes: 0,
83            });
84        entry.rows_in += op.rows_in;
85        entry.rows_out += op.rows_out;
86        entry.batches += op.batches;
87        entry.elapsed_ns += op.elapsed_ns;
88        entry.bytes_h2d += op.bytes_h2d;
89        entry.bytes_d2h += op.bytes_d2h;
90        entry.memory_bytes = entry.memory_bytes.max(op.memory_bytes);
91    }
92    plan.iter()
93        .enumerate()
94        .map(|(id, node)| {
95            let mut op =
96                by_name
97                    .get(node.name.as_str())
98                    .cloned()
99                    .unwrap_or_else(|| OperatorSnapshot {
100                        id: 0,
101                        name: node.name.clone(),
102                        backend: node.backend,
103                        rows_in: 0,
104                        rows_out: 0,
105                        batches: 0,
106                        elapsed_ns: 0,
107                        bytes_h2d: 0,
108                        bytes_d2h: 0,
109                        memory_bytes: 0,
110                    });
111            op.id = id;
112            op.backend = node.backend;
113            op
114        })
115        .collect()
116}
117
118fn render_cell(batch: &RecordBatch, column: usize) -> String {
119    let opts = FormatOptions::default().with_null("");
120    batch
121        .columns()
122        .get(column)
123        .and_then(|c| ArrayFormatter::try_new(c.as_ref(), &opts).ok())
124        .map_or(String::new(), |f| f.value(0).to_string())
125}
126
127/// Profiles every column of `table`: min, max, null count, and P25/P50/P99
128/// (numeric columns only) — one aggregation query per table.
129pub async fn profile_table(
130    session: &OxideSession,
131    table: &str,
132) -> Result<Vec<ColumnProfile>, EngineError> {
133    let schema = session
134        .ctx()
135        .table_provider(table)
136        .await
137        .map_err(EngineError::from)?
138        .schema();
139    let mut selects = vec!["count(*)".to_owned()];
140    // Per column: [count, min?, max?, p25?, p50?, p99?] — track each column's
141    // slot layout so the result row can be unpacked positionally.
142    let mut layout = Vec::new();
143    for field in schema.fields() {
144        let name = field.name();
145        let quoted = format!("\"{}\"", name.replace('"', "\"\""));
146        let numeric = matches!(field.data_type(), DataType::Int64 | DataType::Float64);
147        let orderable = numeric || matches!(field.data_type(), DataType::Utf8);
148        let first = selects.len();
149        selects.push(format!("count({quoted})"));
150        if orderable {
151            selects.push(format!("min({quoted})"));
152            selects.push(format!("max({quoted})"));
153        }
154        if numeric {
155            for q in ["0.25", "0.5", "0.99"] {
156                selects.push(format!("approx_percentile_cont({quoted}, {q})"));
157            }
158        }
159        layout.push((first, orderable, numeric));
160    }
161    // Identifiers come from the CLI (`--table NAME=PATH`); quote them like the
162    // column names above so a name containing `"` cannot escape the query.
163    let quoted_table = format!("\"{}\"", table.replace('"', "\"\""));
164    let sql = format!("SELECT {} FROM {quoted_table}", selects.join(", "));
165    let batches = session.sql(&sql).await?.collect().await?;
166    let row = batches
167        .iter()
168        .find(|b| b.num_rows() > 0)
169        .ok_or_else(|| EngineError::execution("profile query returned no rows"))?;
170    let total = render_cell(row, 0).parse::<u64>().unwrap_or(0);
171    let mut profiles = Vec::with_capacity(schema.fields().len());
172    for (field, (first, orderable, numeric)) in schema.fields().iter().zip(layout) {
173        let non_null = render_cell(row, first).parse::<u64>().unwrap_or(0);
174        let (min, max) = if orderable {
175            (render_cell(row, first + 1), render_cell(row, first + 2))
176        } else {
177            (String::new(), String::new())
178        };
179        let quantile = |i: usize| {
180            if numeric {
181                render_cell(row, first + 3 + i)
182            } else {
183                String::new()
184            }
185        };
186        profiles.push(ColumnProfile {
187            name: field.name().clone(),
188            data_type: field.data_type().to_string(),
189            min,
190            max,
191            null_count: total.saturating_sub(non_null),
192            p25: quantile(0),
193            p50: quantile(1),
194            p99: quantile(2),
195        });
196    }
197    Ok(profiles)
198}
199
200/// Runs `sql` on `session` to completion and returns the dashboard model for
201/// it: the executed plan with placement tags, the aggregated telemetry the
202/// run produced, and column profiles for `tables`.
203pub async fn query_dashboard(
204    session: &OxideSession,
205    sql: &str,
206    tables: &[String],
207) -> Result<DashboardModel, EngineError> {
208    let plan = session
209        .sql(sql)
210        .await?
211        .create_physical_plan()
212        .await
213        .map_err(EngineError::from)?;
214    collect(Arc::clone(&plan), session.ctx().task_ctx())
215        .await
216        .map_err(EngineError::from)?;
217    let nodes = plan_summary(&plan);
218    session.telemetry().set_plan(nodes.clone());
219    let mut telemetry = session.telemetry().snapshot();
220    telemetry.operators = operators_for_plan(&nodes, &telemetry.operators);
221    telemetry.plan = nodes;
222    let mut profiles = Vec::new();
223    for table in tables {
224        profiles.extend(profile_table(session, table).await?);
225    }
226    Ok(DashboardModel {
227        telemetry,
228        profiles,
229        title: sql.to_owned(),
230    })
231}
232
233#[cfg(test)]
234#[allow(clippy::unwrap_used, clippy::expect_used)]
235mod tests {
236    use arrow::array::{Float64Array, Int64Array};
237    use arrow::datatypes::{Field, Schema};
238
239    use super::*;
240
241    fn session_with_table() -> OxideSession {
242        let session = OxideSession::local_with_target(BackendKind::Cuda).unwrap();
243        let schema = Arc::new(Schema::new(vec![
244            Field::new("k", DataType::Int64, true),
245            Field::new("v", DataType::Float64, true),
246        ]));
247        let batch = RecordBatch::try_new(
248            schema,
249            vec![
250                Arc::new(Int64Array::from(vec![Some(1), Some(2), None, Some(2)])),
251                Arc::new(Float64Array::from(vec![
252                    Some(0.5),
253                    Some(1.5),
254                    Some(2.5),
255                    None,
256                ])),
257            ],
258        )
259        .unwrap();
260        session.ctx().register_batch("t", batch).unwrap();
261        session
262    }
263
264    #[tokio::test]
265    async fn dashboard_reports_plan_tags_and_live_counters() {
266        let session = session_with_table();
267        let model = query_dashboard(
268            &session,
269            "SELECT k, sum(v) FROM t WHERE k >= 1 GROUP BY k",
270            &["t".to_owned()],
271        )
272        .await
273        .unwrap();
274        let names: Vec<&str> = model.plan().iter().map(|n| n.name.as_str()).collect();
275        assert!(names.contains(&"GpuAggregateExec"), "{names:?}");
276        assert!(names.contains(&"GpuFilterExec"), "{names:?}");
277        let agg_row = model
278            .plan()
279            .iter()
280            .position(|n| n.name == "GpuAggregateExec")
281            .unwrap();
282        assert_eq!(model.plan()[agg_row].backend, BackendKind::Cuda);
283        // The run recorded real batches through the telemetry hub.
284        let op = model.operator(agg_row).unwrap();
285        assert_eq!(op.name, "GpuAggregateExec");
286        assert!(op.rows_in > 0 && op.batches > 0, "{op:?}");
287        // Profiles: k and v with correct null counts.
288        assert_eq!(model.profiles.len(), 2);
289        assert_eq!(model.profiles[0].name, "k");
290        assert_eq!(model.profiles[0].null_count, 1);
291        assert_eq!(model.profiles[0].min, "1");
292        assert_eq!(model.profiles[0].max, "2");
293        assert_eq!(model.profiles[1].null_count, 1);
294        assert!(!model.profiles[1].p50.is_empty());
295    }
296}
oxidelake_runtime/dashboard.rs

oxidelake_runtime/
dashboard.rs
