Skip to main content

micromegas_analytics/
metrics_table.rs

1use crate::{
2    measure::Measure, metadata::ProcessMetadata,
3    properties::property_set_jsonb_dictionary_builder::PropertySetJsonbDictionaryBuilder,
4    time::TimeRange,
5};
6use anyhow::{Context, Result};
7use chrono::DateTime;
8use datafusion::arrow::{
9    array::{ArrayBuilder, BinaryDictionaryBuilder, PrimitiveBuilder, StringDictionaryBuilder},
10    datatypes::{
11        DataType, Field, Float64Type, Int32Type, Schema, TimeUnit, TimestampNanosecondType,
12    },
13    record_batch::RecordBatch,
14};
15use std::sync::Arc;
16
17/// Returns the schema for the metrics table.
18pub fn metrics_table_schema() -> Schema {
19    Schema::new(vec![
20        Field::new(
21            "process_id",
22            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
23            false,
24        ),
25        Field::new(
26            "stream_id",
27            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
28            false,
29        ),
30        Field::new(
31            "block_id",
32            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
33            false,
34        ),
35        Field::new(
36            "insert_time",
37            DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
38            false,
39        ),
40        Field::new(
41            "exe",
42            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
43            false,
44        ),
45        Field::new(
46            "username",
47            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
48            false,
49        ),
50        Field::new(
51            "computer",
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            "target",
62            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
63            false,
64        ),
65        Field::new(
66            "name",
67            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
68            false,
69        ),
70        Field::new(
71            "unit",
72            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
73            false,
74        ),
75        Field::new("value", DataType::Float64, false),
76        Field::new(
77            "properties",
78            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)),
79            false,
80        ),
81        Field::new(
82            "process_properties",
83            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)),
84            false,
85        ),
86        // Appended last: the owning process's audience, dictionary-encoded like
87        // `process_id` in this same schema. Non-nullable because both producers of this column
88        // resolve a never-stamped process to `MICROMEGAS_DEFAULT_AUDIENCE` before it gets here
89        // (`blocks_view.rs`'s `data_sql` for the global instances, `metadata::find_process` for
90        // the per-process JIT ones).
91        Field::new(
92            "audience",
93            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
94            false,
95        ),
96    ])
97}
98
99/// A builder for creating a `RecordBatch` of metrics.
100pub struct MetricsRecordBuilder {
101    pub process_ids: StringDictionaryBuilder<Int32Type>,
102    pub stream_ids: StringDictionaryBuilder<Int32Type>,
103    pub block_ids: StringDictionaryBuilder<Int32Type>,
104    pub insert_times: PrimitiveBuilder<TimestampNanosecondType>,
105    pub exes: StringDictionaryBuilder<Int32Type>,
106    pub usernames: StringDictionaryBuilder<Int32Type>,
107    pub computers: StringDictionaryBuilder<Int32Type>,
108    pub times: PrimitiveBuilder<TimestampNanosecondType>,
109    pub targets: StringDictionaryBuilder<Int32Type>,
110    pub names: StringDictionaryBuilder<Int32Type>,
111    pub units: StringDictionaryBuilder<Int32Type>,
112    pub values: PrimitiveBuilder<Float64Type>,
113    pub properties: PropertySetJsonbDictionaryBuilder,
114    pub process_properties: BinaryDictionaryBuilder<Int32Type>,
115    pub audiences: StringDictionaryBuilder<Int32Type>,
116}
117
118impl MetricsRecordBuilder {
119    pub fn with_capacity(capacity: usize) -> Self {
120        Self {
121            process_ids: StringDictionaryBuilder::new(),
122            stream_ids: StringDictionaryBuilder::new(),
123            block_ids: StringDictionaryBuilder::new(),
124            insert_times: PrimitiveBuilder::with_capacity(capacity),
125            exes: StringDictionaryBuilder::new(),
126            usernames: StringDictionaryBuilder::new(),
127            computers: StringDictionaryBuilder::new(),
128            times: PrimitiveBuilder::with_capacity(capacity),
129            targets: StringDictionaryBuilder::new(),
130            names: StringDictionaryBuilder::new(),
131            units: StringDictionaryBuilder::new(),
132            values: PrimitiveBuilder::with_capacity(capacity),
133            properties: PropertySetJsonbDictionaryBuilder::new(capacity),
134            process_properties: BinaryDictionaryBuilder::new(),
135            audiences: StringDictionaryBuilder::new(),
136        }
137    }
138
139    pub fn len(&self) -> i64 {
140        self.times.len() as i64
141    }
142
143    pub fn is_empty(&self) -> bool {
144        self.times.len() == 0
145    }
146
147    pub fn get_time_range(&self) -> Option<TimeRange> {
148        if self.is_empty() {
149            return None;
150        }
151        // assuming that the events are in order
152        let slice = self.times.values_slice();
153        Some(TimeRange::new(
154            DateTime::from_timestamp_nanos(slice[0]),
155            DateTime::from_timestamp_nanos(slice[slice.len() - 1]),
156        ))
157    }
158
159    pub fn append(&mut self, row: &Measure) -> Result<()> {
160        self.process_ids
161            .append(format!("{}", row.process.process_id))?;
162        self.stream_ids.append(&*row.stream_id)?;
163        self.block_ids.append(&*row.block_id)?;
164        self.insert_times.append_value(row.insert_time);
165        self.exes.append(&row.process.exe)?;
166        self.usernames.append(&row.process.username)?;
167        self.computers.append(&row.process.computer)?;
168        self.times.append_value(row.time);
169        self.targets.append(row.target)?;
170        self.names.append(row.name)?;
171        self.units.append(row.unit)?;
172        self.values.append_value(row.value);
173        self.properties.append_property_set(&row.properties)?;
174        self.process_properties.append(&*row.process.properties)?;
175        self.audiences.append(&*row.process.audience)?;
176        Ok(())
177    }
178
179    /// Append only per-entry variable data (optimized for batch processing)
180    pub fn append_entry_only(&mut self, row: &Measure) -> Result<()> {
181        // Only append fields that truly vary per metrics entry
182        self.times.append_value(row.time);
183        self.targets.append(row.target)?;
184        self.names.append(row.name)?;
185        self.units.append(row.unit)?;
186        self.values.append_value(row.value);
187        self.properties.append_property_set(&row.properties)?;
188        Ok(())
189    }
190
191    /// Batch fill all constant columns for all entries in block
192    pub fn fill_constant_columns(
193        &mut self,
194        process: &ProcessMetadata,
195        stream_id: &str,
196        block_id: &str,
197        insert_time: i64,
198        entry_count: usize,
199    ) -> Result<()> {
200        let process_id_str = format!("{}", process.process_id);
201
202        // For PrimitiveBuilder (insert_times): use append_slice for better performance
203        let insert_times_slice = vec![insert_time; entry_count];
204        self.insert_times.append_slice(&insert_times_slice);
205
206        // For BinaryDictionaryBuilder (process_properties): use append_n for same value
207        self.process_properties
208            .append_n(&**process.properties, entry_count)?;
209
210        // For StringDictionaryBuilder: use append_n for same values (optimal for constant data)
211        self.process_ids.append_n(&process_id_str, entry_count)?;
212        self.stream_ids.append_n(stream_id, entry_count)?;
213        self.block_ids.append_n(block_id, entry_count)?;
214        self.exes.append_n(&process.exe, entry_count)?;
215        self.usernames.append_n(&process.username, entry_count)?;
216        self.computers.append_n(&process.computer, entry_count)?;
217        self.audiences.append_n(&*process.audience, entry_count)?;
218
219        Ok(())
220    }
221
222    pub fn finish(mut self) -> Result<RecordBatch> {
223        RecordBatch::try_new(
224            Arc::new(metrics_table_schema()),
225            vec![
226                Arc::new(self.process_ids.finish()),
227                Arc::new(self.stream_ids.finish()),
228                Arc::new(self.block_ids.finish()),
229                Arc::new(self.insert_times.finish().with_timezone_utc()),
230                Arc::new(self.exes.finish()),
231                Arc::new(self.usernames.finish()),
232                Arc::new(self.computers.finish()),
233                Arc::new(self.times.finish().with_timezone_utc()),
234                Arc::new(self.targets.finish()),
235                Arc::new(self.names.finish()),
236                Arc::new(self.units.finish()),
237                Arc::new(self.values.finish()),
238                Arc::new(self.properties.finish()?),
239                Arc::new(self.process_properties.finish()),
240                Arc::new(self.audiences.finish()),
241            ],
242        )
243        .with_context(|| "building record batch")
244    }
245}