micromegas_analytics/
async_events_table.rs

1use std::sync::Arc;
2
3use anyhow::{Context, Result};
4use chrono::DateTime;
5use datafusion::arrow::array::ArrayBuilder;
6use datafusion::arrow::array::PrimitiveBuilder;
7use datafusion::arrow::array::StringDictionaryBuilder;
8use datafusion::arrow::datatypes::DataType;
9use datafusion::arrow::datatypes::Field;
10use datafusion::arrow::datatypes::Int32Type;
11use datafusion::arrow::datatypes::Int64Type;
12use datafusion::arrow::datatypes::Schema;
13use datafusion::arrow::datatypes::TimeUnit;
14use datafusion::arrow::datatypes::TimestampNanosecondType;
15use datafusion::arrow::datatypes::UInt32Type;
16use datafusion::arrow::record_batch::RecordBatch;
17
18use crate::time::TimeRange;
19
20/// Represents a single async span event record.
21/// Optimized for high-frequency data - process info can be joined when needed.
22///
23/// `name`/`filename`/`target` borrow the per-block parse arena; the record is
24/// appended to Arrow (which copies the strings) within the parse callback.
25#[derive(Debug, Clone)]
26pub struct AsyncEventRecord<'a> {
27    pub stream_id: Arc<String>,
28    pub block_id: Arc<String>,
29    pub time: i64,
30    pub event_type: &'static str,
31    pub span_id: i64,
32    pub parent_span_id: i64,
33    pub depth: u32,
34    pub hash: u32,
35    pub name: &'a str,
36    pub filename: &'a str,
37    pub target: &'a str,
38    pub line: u32,
39}
40
41/// Returns the schema for the async events table.
42/// Optimized for high-frequency data - excludes process info that can be joined.
43pub fn async_events_table_schema() -> Schema {
44    Schema::new(vec![
45        Field::new(
46            "stream_id",
47            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
48            false,
49        ),
50        Field::new(
51            "block_id",
52            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
53            false,
54        ),
55        Field::new(
56            "time",
57            DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
58            false,
59        ),
60        Field::new(
61            "event_type",
62            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
63            false,
64        ),
65        Field::new("span_id", DataType::Int64, false),
66        Field::new("parent_span_id", DataType::Int64, false),
67        Field::new("depth", DataType::UInt32, false),
68        Field::new("hash", DataType::UInt32, false),
69        Field::new(
70            "name",
71            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
72            false,
73        ),
74        Field::new(
75            "filename",
76            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
77            false,
78        ),
79        Field::new(
80            "target",
81            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
82            false,
83        ),
84        Field::new("line", DataType::UInt32, false),
85    ])
86}
87
88/// A builder for creating a `RecordBatch` of async event records.
89pub struct AsyncEventRecordBuilder {
90    stream_ids: StringDictionaryBuilder<Int32Type>,
91    block_ids: StringDictionaryBuilder<Int32Type>,
92    times: PrimitiveBuilder<TimestampNanosecondType>,
93    event_types: StringDictionaryBuilder<Int32Type>,
94    span_ids: PrimitiveBuilder<Int64Type>,
95    parent_span_ids: PrimitiveBuilder<Int64Type>,
96    depths: PrimitiveBuilder<UInt32Type>,
97    hashes: PrimitiveBuilder<UInt32Type>,
98    names: StringDictionaryBuilder<Int32Type>,
99    filenames: StringDictionaryBuilder<Int32Type>,
100    targets: StringDictionaryBuilder<Int32Type>,
101    lines: PrimitiveBuilder<UInt32Type>,
102}
103
104impl AsyncEventRecordBuilder {
105    pub fn with_capacity(capacity: usize) -> Self {
106        Self {
107            stream_ids: StringDictionaryBuilder::new(),
108            block_ids: StringDictionaryBuilder::new(),
109            times: PrimitiveBuilder::with_capacity(capacity),
110            event_types: StringDictionaryBuilder::new(),
111            span_ids: PrimitiveBuilder::with_capacity(capacity),
112            parent_span_ids: PrimitiveBuilder::with_capacity(capacity),
113            depths: PrimitiveBuilder::with_capacity(capacity),
114            hashes: PrimitiveBuilder::with_capacity(capacity),
115            names: StringDictionaryBuilder::new(),
116            filenames: StringDictionaryBuilder::new(),
117            targets: StringDictionaryBuilder::new(),
118            lines: PrimitiveBuilder::with_capacity(capacity),
119        }
120    }
121
122    pub fn get_time_range(&self) -> Option<TimeRange> {
123        if self.is_empty() {
124            return None;
125        }
126        // assuming that the events are in order
127        let slice = self.times.values_slice();
128        Some(TimeRange::new(
129            DateTime::from_timestamp_nanos(slice[0]),
130            DateTime::from_timestamp_nanos(slice[slice.len() - 1]),
131        ))
132    }
133
134    pub fn len(&self) -> i64 {
135        self.times.len() as i64
136    }
137
138    pub fn is_empty(&self) -> bool {
139        self.times.len() == 0
140    }
141
142    pub fn append(&mut self, record: &AsyncEventRecord<'_>) -> Result<()> {
143        self.stream_ids.append(&*record.stream_id)?;
144        self.block_ids.append(&*record.block_id)?;
145        self.times.append_value(record.time);
146        self.event_types.append(record.event_type)?;
147        self.span_ids.append_value(record.span_id);
148        self.parent_span_ids.append_value(record.parent_span_id);
149        self.depths.append_value(record.depth);
150        self.hashes.append_value(record.hash);
151        self.names.append(record.name)?;
152        self.filenames.append(record.filename)?;
153        self.targets.append(record.target)?;
154        self.lines.append_value(record.line);
155        Ok(())
156    }
157
158    pub fn finish(mut self) -> Result<RecordBatch> {
159        RecordBatch::try_new(
160            Arc::new(async_events_table_schema()),
161            vec![
162                Arc::new(self.stream_ids.finish()),
163                Arc::new(self.block_ids.finish()),
164                Arc::new(self.times.finish().with_timezone_utc()),
165                Arc::new(self.event_types.finish()),
166                Arc::new(self.span_ids.finish()),
167                Arc::new(self.parent_span_ids.finish()),
168                Arc::new(self.depths.finish()),
169                Arc::new(self.hashes.finish()),
170                Arc::new(self.names.finish()),
171                Arc::new(self.filenames.finish()),
172                Arc::new(self.targets.finish()),
173                Arc::new(self.lines.finish()),
174            ],
175        )
176        .with_context(|| "building record batch")
177    }
178}
micromegas_analytics/async_events_table.rs

micromegas_analytics/
async_events_table.rs
