linkrs_query/executor/explain/
execution_stats_context.rs

1//! Execution Statistics Context
2//!
3//! Used to collect execution statistics for all nodes during EXPLAIN ANALYZE and PROFILE operations.
4//! Provides global statistics management and per-node statistics collection.
5
6use std::collections::HashMap;
7use std::time::Instant;
8
9use parking_lot::Mutex;
10
11use crate::executor::base::ExecutorStats;
12use linkrs_metrics::utils::micros_to_millis;
13
14/// Node-level execution statistics
15#[derive(Debug, Clone)]
16pub struct NodeExecutionStats {
17    pub node_id: i64,
18    pub executor_stats: ExecutorStats,
19    pub startup_time_us: u64,
20}
21
22impl NodeExecutionStats {
23    pub fn new(node_id: i64) -> Self {
24        Self {
25            node_id,
26            executor_stats: ExecutorStats::default(),
27            startup_time_us: 0,
28        }
29    }
30
31    pub fn actual_rows(&self) -> usize {
32        self.executor_stats.num_rows
33    }
34
35    pub fn actual_time_us(&self) -> u64 {
36        self.executor_stats.exec_time_us
37    }
38
39    pub fn actual_time_ms(&self) -> f64 {
40        micros_to_millis(self.executor_stats.exec_time_us)
41    }
42
43    pub fn memory_used(&self) -> usize {
44        self.executor_stats.memory_peak
45    }
46}
47
48impl Default for NodeExecutionStats {
49    fn default() -> Self {
50        Self::new(0)
51    }
52}
53
54/// Global execution statistics
55#[derive(Debug, Clone, Default)]
56pub struct GlobalExecutionStats {
57    pub planning_time_us: u64,
58    pub execution_time_us: u64,
59    pub total_rows: usize,
60    pub peak_memory: usize,
61}
62
63impl GlobalExecutionStats {
64    pub fn planning_time_ms(&self) -> f64 {
65        micros_to_millis(self.planning_time_us)
66    }
67
68    pub fn execution_time_ms(&self) -> f64 {
69        micros_to_millis(self.execution_time_us)
70    }
71}
72
73/// Execution statistics context
74///
75/// Manages statistics collection for all plan nodes during query execution.
76/// Used by EXPLAIN ANALYZE and PROFILE to gather actual execution data.
77pub struct ExecutionStatsContext {
78    node_stats: Mutex<HashMap<i64, NodeExecutionStats>>,
79    global_stats: Mutex<GlobalExecutionStats>,
80    start_time: Instant,
81}
82
83impl ExecutionStatsContext {
84    pub fn new() -> Self {
85        Self {
86            node_stats: Mutex::new(HashMap::new()),
87            global_stats: Mutex::new(GlobalExecutionStats::default()),
88            start_time: Instant::now(),
89        }
90    }
91
92    pub fn with_planning_time(planning_time_ms: f64) -> Self {
93        let ctx = Self::new();
94        ctx.global_stats.lock().planning_time_us = (planning_time_ms * 1000.0) as u64;
95        ctx
96    }
97
98    pub fn on_node_start(&self, node_id: i64) {
99        let mut stats = self.node_stats.lock();
100        stats
101            .entry(node_id)
102            .or_insert_with(|| NodeExecutionStats::new(node_id));
103    }
104
105    pub fn on_node_complete(&self, node_id: i64, executor_stats: ExecutorStats) {
106        let mut stats = self.node_stats.lock();
107        let node_stats = NodeExecutionStats {
108            node_id,
109            executor_stats,
110            startup_time_us: 0,
111        };
112        stats.insert(node_id, node_stats);
113    }
114
115    pub fn record_node_rows(&self, node_id: i64, rows: usize) {
116        let mut stats = self.node_stats.lock();
117        if let Some(s) = stats.get_mut(&node_id) {
118            s.executor_stats.num_rows = rows;
119        }
120    }
121
122    pub fn record_node_time(&self, node_id: i64, time_us: u64) {
123        let mut stats = self.node_stats.lock();
124        if let Some(s) = stats.get_mut(&node_id) {
125            s.executor_stats.exec_time_us = time_us;
126        }
127    }
128
129    pub fn record_startup_time(&self, node_id: i64, startup_time_us: u64) {
130        let mut stats = self.node_stats.lock();
131        if let Some(s) = stats.get_mut(&node_id) {
132            s.startup_time_us = startup_time_us;
133        }
134    }
135
136    pub fn record_global_execution_time(&self, time_us: u64) {
137        self.global_stats.lock().execution_time_us = time_us;
138    }
139
140    pub fn collect_stats(&self) -> HashMap<i64, NodeExecutionStats> {
141        self.node_stats.lock().clone()
142    }
143
144    pub fn get_node_stats(&self, node_id: i64) -> Option<NodeExecutionStats> {
145        self.node_stats.lock().get(&node_id).cloned()
146    }
147
148    pub fn get_global_stats(&self) -> GlobalExecutionStats {
149        self.global_stats.lock().clone()
150    }
151
152    pub fn total_elapsed_ms(&self) -> f64 {
153        self.start_time.elapsed().as_micros() as f64 / 1000.0
154    }
155}
156
157impl Default for ExecutionStatsContext {
158    fn default() -> Self {
159        Self::new()
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn test_execution_stats_context() {
169        let ctx = ExecutionStatsContext::new();
170
171        ctx.on_node_start(1);
172        let exec_stats = ExecutorStats {
173            num_rows: 100,
174            exec_time_us: 5500,
175            ..ExecutorStats::default()
176        };
177        ctx.on_node_complete(1, exec_stats);
178
179        let collected = ctx.collect_stats();
180        assert_eq!(collected.get(&1).unwrap().actual_rows(), 100);
181        assert!((collected.get(&1).unwrap().actual_time_ms() - 5.5).abs() < 0.001);
182    }
183}
linkrs_query/executor/explain/execution_stats_context.rs

linkrs_query/executor/explain/
execution_stats_context.rs
