micromegas_analytics/lakehouse/
thread_spans_view.rs

1use super::{
2    blocks_view::BlocksView,
3    dataframe_time_bounds::{DataFrameTimeBounds, NamedColumnsTimeBounds},
4    jit_partitions::{
5        JitPartitionConfig, generate_stream_jit_partitions, is_jit_partition_up_to_date,
6    },
7    lakehouse_context::LakehouseContext,
8    partition_cache::PartitionCache,
9    partition_source_data::{SourceDataBlocksInMemory, hash_to_object_count},
10    view::{PartitionSpec, ScanSortColumn, View, ViewMetadata},
11    view_factory::{ViewFactory, ViewMaker},
12};
13use crate::{
14    call_tree::make_call_tree,
15    lakehouse::write_partition::{PartitionRowSet, write_partition_from_rows},
16    metadata::{find_process_with_latest_timing, find_stream_from_view},
17    response_writer::ResponseWriter,
18    span_table::{SpanRecordBuilder, get_spans_schema},
19    time::{ConvertTicks, TimeRange, datetime_to_scalar, make_time_converter_from_latest_timing},
20};
21use anyhow::{Context, Result};
22use async_trait::async_trait;
23use chrono::{DateTime, Utc};
24use datafusion::logical_expr::{BinaryExpr, Expr, Operator};
25use datafusion::{arrow::datatypes::Schema, logical_expr::expr_fn::col};
26use micromegas_ingestion::data_lake_connection::DataLakeConnection;
27use micromegas_telemetry::{blob_storage::BlobStorage, types::block::BlockMetadata};
28use micromegas_tracing::prelude::*;
29use std::sync::Arc;
30use uuid::Uuid;
31
32const VIEW_SET_NAME: &str = "thread_spans";
33const SCHEMA_VERSION: u8 = 1;
34lazy_static::lazy_static! {
35    static ref MIN_TIME_COLUMN: Arc<String> = Arc::new( String::from("begin"));
36    static ref MAX_TIME_COLUMN: Arc<String> = Arc::new( String::from("end"));
37}
38
39/// A `ViewMaker` for creating `ThreadSpansView` instances.
40#[derive(Debug)]
41pub struct ThreadSpansViewMaker {
42    view_factory: Arc<ViewFactory>,
43}
44
45impl ThreadSpansViewMaker {
46    pub fn new(view_factory: Arc<ViewFactory>) -> Self {
47        Self { view_factory }
48    }
49}
50
51impl ViewMaker for ThreadSpansViewMaker {
52    fn make_view(&self, stream_id: &str) -> Result<Arc<dyn View>> {
53        Ok(Arc::new(ThreadSpansView::new(
54            stream_id,
55            self.view_factory.clone(),
56        )?))
57    }
58
59    fn get_schema_hash(&self) -> Vec<u8> {
60        vec![SCHEMA_VERSION]
61    }
62
63    fn get_schema(&self) -> Arc<Schema> {
64        Arc::new(get_spans_schema())
65    }
66}
67
68/// A view of thread spans.
69#[derive(Debug)]
70pub struct ThreadSpansView {
71    view_set_name: Arc<String>,
72    view_instance_id: Arc<String>,
73    stream_id: sqlx::types::Uuid,
74    view_factory: Arc<ViewFactory>,
75}
76
77impl ThreadSpansView {
78    pub fn new(view_instance_id: &str, view_factory: Arc<ViewFactory>) -> Result<Self> {
79        if view_instance_id == "global" {
80            anyhow::bail!("the global view is not implemented for thread spans");
81        }
82
83        Ok(Self {
84            view_set_name: Arc::new(String::from(VIEW_SET_NAME)),
85            view_instance_id: Arc::new(String::from(view_instance_id)),
86            stream_id: Uuid::parse_str(view_instance_id).with_context(|| "Uuid::parse_str")?,
87            view_factory,
88        })
89    }
90}
91
92#[span_fn]
93async fn append_call_tree(
94    record_builder: &mut SpanRecordBuilder,
95    convert_ticks: &ConvertTicks,
96    blocks: &[BlockMetadata],
97    blob_storage: Arc<BlobStorage>,
98    stream: &crate::metadata::StreamMetadata,
99) -> Result<()> {
100    let call_tree = make_call_tree(
101        blocks,
102        convert_ticks.delta_ticks_to_ns(blocks[0].begin_ticks),
103        convert_ticks.delta_ticks_to_ns(blocks[blocks.len() - 1].end_ticks),
104        None,
105        blob_storage,
106        convert_ticks.clone(),
107        stream,
108    )
109    .await
110    .with_context(|| "make_call_tree")?;
111    record_builder
112        .append_call_tree(&call_tree)
113        .with_context(|| "adding call tree to span record builder")?;
114    Ok(())
115}
116
117/// Writes a partition from a set of blocks.
118#[span_fn]
119async fn write_partition(
120    lake: Arc<DataLakeConnection>,
121    view_meta: ViewMetadata,
122    schema: Arc<Schema>,
123    convert_ticks: &ConvertTicks,
124    spec: &SourceDataBlocksInMemory,
125) -> Result<()> {
126    let nb_events = hash_to_object_count(&spec.block_ids_hash)? as usize;
127    info!("nb_events: {nb_events}");
128    if spec.blocks.is_empty() {
129        anyhow::bail!("empty partition spec");
130    }
131    // for jit partitions, we assume that the blocks were registered in order
132    // since they are built based on begin_ticks, not insert_time
133    let min_insert_time = spec.blocks[0].block.insert_time;
134    let max_insert_time = spec.blocks[spec.blocks.len() - 1].block.insert_time;
135
136    let (tx, rx) = tokio::sync::mpsc::channel(1);
137    let null_response_writer = Arc::new(ResponseWriter::new(None));
138    let join_handle = spawn_with_context(write_partition_from_rows(
139        lake.clone(),
140        view_meta,
141        schema,
142        TimeRange::new(min_insert_time, max_insert_time),
143        spec.block_ids_hash.clone(),
144        None,
145        rx,
146        null_response_writer,
147    ));
148
149    let build_result: Result<PartitionRowSet> = async {
150        let mut record_builder = SpanRecordBuilder::with_capacity(nb_events / 2);
151        let mut blocks_to_process = vec![];
152        let mut last_end = None;
153        for block in &spec.blocks {
154            if block.block.begin_ticks == last_end.unwrap_or(block.block.begin_ticks) {
155                last_end = Some(block.block.end_ticks);
156                blocks_to_process.push(block.block.clone());
157            } else {
158                append_call_tree(
159                    &mut record_builder,
160                    convert_ticks,
161                    &blocks_to_process,
162                    lake.blob_storage.clone(),
163                    &block.stream,
164                )
165                .await?;
166                last_end = Some(block.block.end_ticks);
167                blocks_to_process = vec![block.block.clone()];
168            }
169        }
170        if !blocks_to_process.is_empty() {
171            append_call_tree(
172                &mut record_builder,
173                convert_ticks,
174                &blocks_to_process,
175                lake.blob_storage.clone(),
176                &spec.blocks[0].stream,
177            )
178            .await?;
179        }
180        let min_time_row = convert_ticks.delta_ticks_to_time(spec.blocks[0].block.begin_ticks);
181        let max_time_row =
182            convert_ticks.delta_ticks_to_time(spec.blocks[spec.blocks.len() - 1].block.end_ticks);
183        let rows = record_builder
184            .finish()
185            .with_context(|| "record_builder.finish()")?;
186        info!("writing {} rows", rows.num_rows());
187        Ok(PartitionRowSet {
188            rows_time_range: TimeRange::new(min_time_row, max_time_row),
189            rows,
190        })
191    }
192    .await;
193
194    match build_result {
195        Ok(row_set) => {
196            tx.send(Ok(row_set)).await?;
197            drop(tx);
198            join_handle.await??;
199            Ok(())
200        }
201        Err(e) => {
202            warn!(
203                "aborting thread-spans partition write for block {:?}: {e:?}",
204                spec.block_ids_hash
205            );
206            let _ = tx
207                .send(Err(anyhow::anyhow!("thread-spans build aborted")))
208                .await;
209            drop(tx);
210            match join_handle.await {
211                Ok(Ok(())) => {}
212                Ok(Err(writer_err)) => {
213                    debug!("thread-spans writer task error during abort: {writer_err:?}");
214                }
215                Err(join_err) => {
216                    warn!("thread-spans writer task panicked during abort: {join_err:?}");
217                }
218            }
219            Err(e)
220        }
221    }
222}
223/// Rebuilds the partition if it's missing or out of date.
224#[span_fn]
225async fn update_partition(
226    lake: Arc<DataLakeConnection>,
227    view_meta: ViewMetadata,
228    schema: Arc<Schema>,
229    convert_ticks: &ConvertTicks,
230    spec: &SourceDataBlocksInMemory,
231) -> Result<()> {
232    if is_jit_partition_up_to_date(&lake.db_pool, view_meta.clone(), spec).await? {
233        return Ok(());
234    }
235    write_partition(lake, view_meta, schema, convert_ticks, spec)
236        .await
237        .with_context(|| "write_partition")?;
238
239    Ok(())
240}
241
242#[async_trait]
243impl View for ThreadSpansView {
244    fn get_view_set_name(&self) -> Arc<String> {
245        self.view_set_name.clone()
246    }
247
248    fn get_view_instance_id(&self) -> Arc<String> {
249        self.view_instance_id.clone()
250    }
251
252    async fn make_batch_partition_spec(
253        &self,
254        _lakehouse: Arc<LakehouseContext>,
255        _existing_partitions: Arc<PartitionCache>,
256        _insert_range: TimeRange,
257    ) -> Result<Arc<dyn PartitionSpec>> {
258        anyhow::bail!("not implemented")
259    }
260
261    fn get_file_schema_hash(&self) -> Vec<u8> {
262        vec![SCHEMA_VERSION]
263    }
264
265    fn get_file_schema(&self) -> Arc<Schema> {
266        Arc::new(get_spans_schema())
267    }
268
269    #[span_fn]
270    async fn jit_update(
271        &self,
272        lakehouse: Arc<LakehouseContext>,
273        query_range: Option<TimeRange>,
274    ) -> Result<()> {
275        let Some(query_range) = query_range else {
276            anyhow::bail!("query range mandatory for thread spans view");
277        };
278        let stream = Arc::new(
279            find_stream_from_view(
280                lakehouse.clone(),
281                self.view_factory.clone(),
282                &self.stream_id,
283                None,
284            )
285            .await
286            .with_context(|| "find_stream_from_view")?,
287        );
288        let (process, last_block_end_ticks, last_block_end_time) = find_process_with_latest_timing(
289            lakehouse.clone(),
290            self.view_factory.clone(),
291            &stream.process_id,
292            None,
293        )
294        .await
295        .with_context(|| "find_process_with_latest_timing")?;
296        let process = Arc::new(process);
297        let convert_ticks = make_time_converter_from_latest_timing(
298            &process,
299            last_block_end_ticks,
300            last_block_end_time,
301        )
302        .with_context(|| "make_time_converter_from_latest_timing")?;
303        let blocks_view = BlocksView::new()?;
304        let partitions = generate_stream_jit_partitions(
305            &JitPartitionConfig::default(),
306            lakehouse.clone(),
307            &blocks_view,
308            &query_range,
309            stream.clone(),
310            process.clone(),
311        )
312        .await
313        .with_context(|| "generate_stream_jit_partitions")?;
314        for part in &partitions {
315            update_partition(
316                lakehouse.lake().clone(),
317                ViewMetadata {
318                    view_set_name: self.get_view_set_name(),
319                    view_instance_id: self.get_view_instance_id(),
320                    file_schema_hash: self.get_file_schema_hash(),
321                },
322                self.get_file_schema(),
323                &convert_ticks,
324                part,
325            )
326            .await
327            .with_context(|| "update_partition")?;
328        }
329        Ok(())
330    }
331
332    fn make_time_filter(&self, begin: DateTime<Utc>, end: DateTime<Utc>) -> Result<Vec<Expr>> {
333        Ok(vec![
334            Expr::BinaryExpr(BinaryExpr::new(
335                col("begin").into(),
336                Operator::LtEq,
337                Expr::Literal(datetime_to_scalar(end), None).into(),
338            )),
339            Expr::BinaryExpr(BinaryExpr::new(
340                col("end").into(),
341                Operator::GtEq,
342                Expr::Literal(datetime_to_scalar(begin), None).into(),
343            )),
344        ])
345    }
346
347    fn get_time_bounds(&self) -> Arc<dyn DataFrameTimeBounds> {
348        Arc::new(NamedColumnsTimeBounds::new(
349            MIN_TIME_COLUMN.clone(),
350            MAX_TIME_COLUMN.clone(),
351        ))
352    }
353
354    fn get_update_group(&self) -> Option<i32> {
355        None
356    }
357
358    fn get_scan_output_ordering(&self) -> Vec<ScanSortColumn> {
359        vec![ScanSortColumn {
360            column: MIN_TIME_COLUMN.clone(),
361            descending: false,
362        }]
363    }
364}
micromegas_analytics/lakehouse/thread_spans_view.rs

micromegas_analytics/lakehouse/
thread_spans_view.rs
