micromegas_analytics/lakehouse/
process_spans_table_function.rs

1use super::{
2    lakehouse_context::LakehouseContext, partition_cache::QueryPartitionProvider,
3    process_streams::get_process_thread_list, session_configurator::NoOpSessionConfigurator,
4    view_factory::ViewFactory,
5};
6use crate::{dfext::expressions::exp_to_string, span_table::get_spans_schema, time::TimeRange};
7use async_stream::try_stream;
8use datafusion::{
9    arrow::{
10        array::{ArrayRef, RecordBatch, StringDictionaryBuilder},
11        datatypes::{DataType, Field, Int16Type, Schema, SchemaRef},
12    },
13    catalog::{Session, TableFunctionArgs, TableFunctionImpl, TableProvider},
14    common::{Result as DFResult, plan_err},
15    execution::{SendableRecordBatchStream, TaskContext},
16    logical_expr::{Expr, TableType},
17    physical_expr::EquivalenceProperties,
18    physical_plan::{
19        DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
20        execution_plan::{Boundedness, EmissionType},
21        limit::GlobalLimitExec,
22        projection::ProjectionExec,
23        stream::RecordBatchStreamAdapter,
24    },
25};
26use futures::{StreamExt, TryStreamExt};
27use micromegas_tracing::prelude::*;
28use std::{
29    fmt::{self, Debug, Formatter},
30    sync::Arc,
31};
32
33/// Span types to include in the output
34#[derive(Debug, Clone, Copy)]
35pub enum SpanTypes {
36    Thread,
37    Async,
38    Both,
39}
40
41fn output_schema() -> SchemaRef {
42    let mut fields = vec![
43        Field::new(
44            "stream_id",
45            DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)),
46            false,
47        ),
48        Field::new(
49            "thread_name",
50            DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)),
51            false,
52        ),
53    ];
54    fields.extend(get_spans_schema().fields.iter().map(|f| f.as_ref().clone()));
55    Arc::new(Schema::new(fields))
56}
57
58fn augment_batch(
59    batch: &RecordBatch,
60    schema: SchemaRef,
61    stream_id: &str,
62    thread_name: &str,
63) -> DFResult<RecordBatch> {
64    let n = batch.num_rows();
65    let mut stream_id_builder = StringDictionaryBuilder::<Int16Type>::new();
66    let mut thread_name_builder = StringDictionaryBuilder::<Int16Type>::new();
67    stream_id_builder.append_values(stream_id, n);
68    thread_name_builder.append_values(thread_name, n);
69    let mut columns: Vec<ArrayRef> = vec![
70        Arc::new(stream_id_builder.finish()),
71        Arc::new(thread_name_builder.finish()),
72    ];
73    columns.extend(batch.columns().iter().cloned());
74    RecordBatch::try_new(schema, columns).map_err(Into::into)
75}
76
77// --- TableFunction ---
78
79#[derive(Debug)]
80pub struct ProcessSpansTableFunction {
81    lakehouse: Arc<LakehouseContext>,
82    view_factory: Arc<ViewFactory>,
83    part_provider: Arc<dyn QueryPartitionProvider>,
84    query_range: Option<TimeRange>,
85}
86
87impl ProcessSpansTableFunction {
88    pub fn new(
89        lakehouse: Arc<LakehouseContext>,
90        view_factory: Arc<ViewFactory>,
91        part_provider: Arc<dyn QueryPartitionProvider>,
92        query_range: Option<TimeRange>,
93    ) -> Self {
94        Self {
95            lakehouse,
96            view_factory,
97            part_provider,
98            query_range,
99        }
100    }
101}
102
103impl TableFunctionImpl for ProcessSpansTableFunction {
104    #[span_fn]
105    fn call_with_args(
106        &self,
107        args: TableFunctionArgs,
108    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
109        let exprs = args.exprs();
110        let arg1 = exprs.first().map(exp_to_string);
111        let Some(Ok(process_id)) = arg1 else {
112            return plan_err!(
113                "First argument to process_spans must be a string (the process ID), given {:?}",
114                arg1
115            );
116        };
117
118        let arg2 = exprs.get(1).map(exp_to_string);
119        let Some(Ok(span_types_str)) = arg2 else {
120            return plan_err!(
121                "Second argument to process_spans must be a string ('thread', 'async', or 'both'), given {:?}",
122                arg2
123            );
124        };
125
126        let span_types = match span_types_str.as_str() {
127            "thread" => SpanTypes::Thread,
128            "async" => SpanTypes::Async,
129            "both" => SpanTypes::Both,
130            _ => {
131                return plan_err!(
132                    "span_types must be 'thread', 'async', or 'both', given: {span_types_str}"
133                );
134            }
135        };
136
137        let schema = output_schema();
138        let execution_plan = Arc::new(ProcessSpansExecutionPlan::new(
139            schema,
140            process_id,
141            span_types,
142            self.query_range,
143            self.lakehouse.clone(),
144            self.view_factory.clone(),
145            self.part_provider.clone(),
146        ));
147
148        Ok(Arc::new(ProcessSpansTableProvider { execution_plan }))
149    }
150}
151
152// --- ExecutionPlan ---
153
154pub struct ProcessSpansExecutionPlan {
155    schema: SchemaRef,
156    process_id: String,
157    span_types: SpanTypes,
158    query_range: Option<TimeRange>,
159    lakehouse: Arc<LakehouseContext>,
160    view_factory: Arc<ViewFactory>,
161    part_provider: Arc<dyn QueryPartitionProvider>,
162    properties: Arc<PlanProperties>,
163}
164
165impl ProcessSpansExecutionPlan {
166    fn new(
167        schema: SchemaRef,
168        process_id: String,
169        span_types: SpanTypes,
170        query_range: Option<TimeRange>,
171        lakehouse: Arc<LakehouseContext>,
172        view_factory: Arc<ViewFactory>,
173        part_provider: Arc<dyn QueryPartitionProvider>,
174    ) -> Self {
175        let properties = PlanProperties::new(
176            EquivalenceProperties::new(schema.clone()),
177            Partitioning::UnknownPartitioning(1),
178            EmissionType::Final,
179            Boundedness::Bounded,
180        );
181        Self {
182            schema,
183            process_id,
184            span_types,
185            query_range,
186            lakehouse,
187            view_factory,
188            part_provider,
189            properties: Arc::new(properties),
190        }
191    }
192}
193
194impl Debug for ProcessSpansExecutionPlan {
195    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
196        f.debug_struct("ProcessSpansExecutionPlan")
197            .field("process_id", &self.process_id)
198            .field("span_types", &self.span_types)
199            .finish()
200    }
201}
202
203impl DisplayAs for ProcessSpansExecutionPlan {
204    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
205        write!(
206            f,
207            "ProcessSpansExecutionPlan: process_id={}, span_types={:?}",
208            self.process_id, self.span_types
209        )
210    }
211}
212
213impl ExecutionPlan for ProcessSpansExecutionPlan {
214    fn name(&self) -> &str {
215        "ProcessSpansExecutionPlan"
216    }
217
218    fn schema(&self) -> SchemaRef {
219        self.schema.clone()
220    }
221
222    fn properties(&self) -> &Arc<PlanProperties> {
223        &self.properties
224    }
225
226    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
227        vec![]
228    }
229
230    fn with_new_children(
231        self: Arc<Self>,
232        _children: Vec<Arc<dyn ExecutionPlan>>,
233    ) -> DFResult<Arc<dyn ExecutionPlan>> {
234        Ok(self)
235    }
236
237    #[span_fn]
238    fn execute(
239        &self,
240        _partition: usize,
241        _context: Arc<TaskContext>,
242    ) -> DFResult<SendableRecordBatchStream> {
243        let schema = self.schema.clone();
244        let stream_schema = schema.clone();
245        let process_id = self.process_id.clone();
246        let span_types = self.span_types;
247        let query_range = self.query_range;
248        let lakehouse = self.lakehouse.clone();
249        let view_factory = self.view_factory.clone();
250        let part_provider = self.part_provider.clone();
251
252        let record_batch_stream = try_stream! {
253            let schema = stream_schema;
254            let ctx = super::query::make_session_context(
255                lakehouse,
256                part_provider,
257                query_range,
258                view_factory,
259                Arc::new(NoOpSessionConfigurator),
260                false,
261            )
262            .await
263            .map_err(|e| datafusion::error::DataFusionError::Execution(
264                format!("Failed to create session context: {e}"),
265            ))?;
266
267            // Thread spans
268            if matches!(span_types, SpanTypes::Thread | SpanTypes::Both) {
269                let threads = get_process_thread_list(&process_id, &ctx)
270                    .await
271                    .map_err(|e| datafusion::error::DataFusionError::Execution(
272                        format!("Failed to get thread list: {e}"),
273                    ))?;
274
275                let max_concurrent = std::thread::available_parallelism()
276                    .map(|n| n.get())
277                    .unwrap_or(4);
278
279                let queries: Vec<(String, String, String)> = threads
280                    .iter()
281                    .map(|(stream_id, _thread_id, display_name)| {
282                        let sql = format!(
283                            "SELECT * FROM view_instance('thread_spans', '{stream_id}')"
284                        );
285                        (stream_id.clone(), display_name.clone(), sql)
286                    })
287                    .collect();
288
289                let stream_results: Vec<(String, String, SendableRecordBatchStream)> =
290                    futures::stream::iter(queries)
291                        .map(|(stream_id, thread_name, sql)| {
292                            let ctx = ctx.clone();
293                            async move {
294                                spawn_with_context(async move {
295                                    let df = ctx.sql(&sql).await?;
296                                    let s = df.execute_stream().await?;
297                                    Ok::<_, anyhow::Error>((stream_id, thread_name, s))
298                                })
299                                .await?
300                            }
301                        })
302                        .buffered(max_concurrent)
303                        .try_collect()
304                        .await
305                        .map_err(|e| datafusion::error::DataFusionError::Execution(
306                            format!("Failed to query thread spans: {e}"),
307                        ))?;
308
309                for (stream_id, thread_name, mut data_stream) in stream_results {
310                    while let Some(batch) = data_stream.try_next().await? {
311                        let augmented = augment_batch(&batch, schema.clone(), &stream_id, &thread_name)?;
312                        yield augmented;
313                    }
314                }
315            }
316
317            // Async spans
318            if matches!(span_types, SpanTypes::Async | SpanTypes::Both) {
319                let async_sql = format!(
320                    "SELECT \
321                        b.span_id as id, \
322                        b.parent_span_id as parent, \
323                        b.depth, \
324                        b.hash, \
325                        b.time as \"begin\", \
326                        e.time as \"end\", \
327                        arrow_cast(e.time, 'Int64') - arrow_cast(b.time, 'Int64') as duration, \
328                        b.name, \
329                        b.target, \
330                        b.filename, \
331                        b.line \
332                    FROM (SELECT * FROM view_instance('async_events', '{process_id}') \
333                          WHERE event_type = 'begin') b \
334                    INNER JOIN (SELECT * FROM view_instance('async_events', '{process_id}') \
335                          WHERE event_type = 'end') e \
336                    ON b.span_id = e.span_id \
337                    WHERE b.time < e.time \
338                    ORDER BY b.time"
339                );
340
341                let df = ctx.sql(&async_sql).await
342                    .map_err(|e| datafusion::error::DataFusionError::Execution(
343                        format!("Failed to query async spans: {e}"),
344                    ))?;
345                let mut async_stream = df.execute_stream().await
346                    .map_err(|e| datafusion::error::DataFusionError::Execution(
347                        format!("Failed to execute async spans stream: {e}"),
348                    ))?;
349
350                while let Some(batch) = async_stream.try_next().await? {
351                    let augmented = augment_batch(&batch, schema.clone(), "", "async")?;
352                    yield augmented;
353                }
354            }
355        };
356
357        Ok(Box::pin(RecordBatchStreamAdapter::new(
358            schema,
359            record_batch_stream,
360        )))
361    }
362}
363
364// --- TableProvider ---
365
366#[derive(Debug)]
367struct ProcessSpansTableProvider {
368    execution_plan: Arc<ProcessSpansExecutionPlan>,
369}
370
371#[async_trait::async_trait]
372impl TableProvider for ProcessSpansTableProvider {
373    fn schema(&self) -> SchemaRef {
374        self.execution_plan.schema()
375    }
376
377    fn table_type(&self) -> TableType {
378        TableType::Base
379    }
380
381    async fn scan(
382        &self,
383        _state: &dyn Session,
384        projection: Option<&Vec<usize>>,
385        _filters: &[Expr],
386        limit: Option<usize>,
387    ) -> DFResult<Arc<dyn ExecutionPlan>> {
388        let mut plan: Arc<dyn ExecutionPlan> = self.execution_plan.clone();
389        if let Some(projection) = projection {
390            let schema = plan.schema();
391            let projected_exprs: Vec<(Arc<dyn datafusion::physical_expr::PhysicalExpr>, String)> =
392                projection
393                    .iter()
394                    .map(|&i| {
395                        let name = schema.field(i).name().clone();
396                        let expr = Arc::new(datafusion::physical_expr::expressions::Column::new(
397                            &name, i,
398                        ))
399                            as Arc<dyn datafusion::physical_expr::PhysicalExpr>;
400                        (expr, name)
401                    })
402                    .collect();
403            plan = Arc::new(ProjectionExec::try_new(projected_exprs, plan)?);
404        }
405        if let Some(fetch) = limit {
406            plan = Arc::new(GlobalLimitExec::new(plan, 0, Some(fetch)));
407        }
408        Ok(plan)
409    }
410}
micromegas_analytics/lakehouse/process_spans_table_function.rs

micromegas_analytics/lakehouse/
process_spans_table_function.rs
