micromegas_analytics/lakehouse/
perfetto_trace_table_function.rs

1use super::{
2    lakehouse_context::LakehouseContext, partition_cache::QueryPartitionProvider,
3    view_factory::ViewFactory,
4};
5use crate::{
6    dfext::expressions::{exp_to_string, exp_to_timestamp},
7    time::TimeRange,
8};
9use datafusion::{
10    arrow::datatypes::{DataType, Field, Schema},
11    catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider},
12    common::plan_err,
13};
14use micromegas_tracing::prelude::*;
15use std::sync::Arc;
16
17/// `PerfettoTraceTableFunction` generates Perfetto trace chunks from process telemetry data.
18///
19/// SQL Interface:
20/// ```sql
21/// SELECT chunk_id, chunk_data
22/// FROM perfetto_trace_chunks(
23///     'process_id',                              -- Process UUID (required)
24///     'span_types',                              -- 'thread', 'async', or 'both' (required)
25///     TIMESTAMP '2024-01-01T00:00:00Z',          -- Start time as UTC timestamp (required)
26///     TIMESTAMP '2024-01-01T01:00:00Z'           -- End time as UTC timestamp (required)
27/// ) ORDER BY chunk_id
28/// ```
29///
30/// Returns a table with schema:
31/// - chunk_id: Int32 - Sequential chunk identifier
32/// - chunk_data: Binary - Binary protobuf TracePacket data
33///
34#[derive(Debug)]
35pub struct PerfettoTraceTableFunction {
36    lakehouse: Arc<LakehouseContext>,
37    view_factory: Arc<ViewFactory>,
38    part_provider: Arc<dyn QueryPartitionProvider>,
39}
40
41impl PerfettoTraceTableFunction {
42    pub fn new(
43        lakehouse: Arc<LakehouseContext>,
44        view_factory: Arc<ViewFactory>,
45        part_provider: Arc<dyn QueryPartitionProvider>,
46    ) -> Self {
47        Self {
48            lakehouse,
49            view_factory,
50            part_provider,
51        }
52    }
53
54    /// Create the output schema for the table function
55    pub fn output_schema() -> Arc<Schema> {
56        Arc::new(Schema::new(vec![
57            Field::new("chunk_id", DataType::Int32, false),
58            Field::new("chunk_data", DataType::Binary, false),
59        ]))
60    }
61}
62
63impl TableFunctionImpl for PerfettoTraceTableFunction {
64    #[span_fn]
65    fn call_with_args(
66        &self,
67        args: TableFunctionArgs,
68    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
69        let exprs = args.exprs();
70        // Parse process_id (arg 1)
71        let arg1 = exprs.first().map(exp_to_string);
72        let Some(Ok(process_id)) = arg1 else {
73            return plan_err!(
74                "First argument to perfetto_trace_chunks must be a string (the process ID), given {:?}",
75                arg1
76            );
77        };
78
79        // Parse span_types (arg 2)
80        let arg2 = exprs.get(1).map(exp_to_string);
81        let Some(Ok(span_types_str)) = arg2 else {
82            return plan_err!(
83                "Second argument to perfetto_trace_chunks must be a string ('thread', 'async', or 'both'), given {:?}",
84                arg2
85            );
86        };
87
88        let span_types = match span_types_str.as_str() {
89            "thread" => SpanTypes::Thread,
90            "async" => SpanTypes::Async,
91            "both" => SpanTypes::Both,
92            _ => {
93                return plan_err!(
94                    "span_types must be 'thread', 'async', or 'both', given: {}",
95                    span_types_str
96                );
97            }
98        };
99
100        // Parse start_time (arg 3) - expecting a timestamp expression
101        let arg3 = exprs.get(2).map(exp_to_timestamp);
102        let Some(Ok(start_time)) = arg3 else {
103            return plan_err!(
104                "Third argument to perfetto_trace_chunks must be a timestamp (start time), given {:?}",
105                arg3
106            );
107        };
108
109        // Parse end_time (arg 4) - expecting a timestamp expression
110        let arg4 = exprs.get(3).map(exp_to_timestamp);
111        let Some(Ok(end_time)) = arg4 else {
112            return plan_err!(
113                "Fourth argument to perfetto_trace_chunks must be a timestamp (end time), given {:?}",
114                arg4
115            );
116        };
117
118        // Create time range from parsed timestamps
119        let time_range = TimeRange {
120            begin: start_time,
121            end: end_time,
122        };
123
124        // Create the execution plan that will generate the trace chunks
125        let execution_plan = Arc::new(PerfettoTraceExecutionPlan::new(
126            Self::output_schema(),
127            process_id,
128            span_types,
129            time_range,
130            self.lakehouse.clone(),
131            self.view_factory.clone(),
132            self.part_provider.clone(),
133        ));
134
135        // Wrap it in a TableProvider
136        Ok(Arc::new(PerfettoTraceTableProvider::new(execution_plan)))
137    }
138}
139
140// Import the execution plan
141use super::perfetto_trace_execution_plan::{
142    PerfettoTraceExecutionPlan, PerfettoTraceTableProvider, SpanTypes,
143};
micromegas_analytics/lakehouse/perfetto_trace_table_function.rs

micromegas_analytics/lakehouse/
perfetto_trace_table_function.rs
