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    ])
87}
88
89/// A builder for creating a `RecordBatch` of metrics.
90pub struct MetricsRecordBuilder {
91    pub process_ids: StringDictionaryBuilder<Int32Type>,
92    pub stream_ids: StringDictionaryBuilder<Int32Type>,
93    pub block_ids: StringDictionaryBuilder<Int32Type>,
94    pub insert_times: PrimitiveBuilder<TimestampNanosecondType>,
95    pub exes: StringDictionaryBuilder<Int32Type>,
96    pub usernames: StringDictionaryBuilder<Int32Type>,
97    pub computers: StringDictionaryBuilder<Int32Type>,
98    pub times: PrimitiveBuilder<TimestampNanosecondType>,
99    pub targets: StringDictionaryBuilder<Int32Type>,
100    pub names: StringDictionaryBuilder<Int32Type>,
101    pub units: StringDictionaryBuilder<Int32Type>,
102    pub values: PrimitiveBuilder<Float64Type>,
103    pub properties: PropertySetJsonbDictionaryBuilder,
104    pub process_properties: BinaryDictionaryBuilder<Int32Type>,
105}
106
107impl MetricsRecordBuilder {
108    pub fn with_capacity(capacity: usize) -> Self {
109        Self {
110            process_ids: StringDictionaryBuilder::new(),
111            stream_ids: StringDictionaryBuilder::new(),
112            block_ids: StringDictionaryBuilder::new(),
113            insert_times: PrimitiveBuilder::with_capacity(capacity),
114            exes: StringDictionaryBuilder::new(),
115            usernames: StringDictionaryBuilder::new(),
116            computers: StringDictionaryBuilder::new(),
117            times: PrimitiveBuilder::with_capacity(capacity),
118            targets: StringDictionaryBuilder::new(),
119            names: StringDictionaryBuilder::new(),
120            units: StringDictionaryBuilder::new(),
121            values: PrimitiveBuilder::with_capacity(capacity),
122            properties: PropertySetJsonbDictionaryBuilder::new(capacity),
123            process_properties: BinaryDictionaryBuilder::new(),
124        }
125    }
126
127    pub fn len(&self) -> i64 {
128        self.times.len() as i64
129    }
130
131    pub fn is_empty(&self) -> bool {
132        self.times.len() == 0
133    }
134
135    pub fn get_time_range(&self) -> Option<TimeRange> {
136        if self.is_empty() {
137            return None;
138        }
139        // assuming that the events are in order
140        let slice = self.times.values_slice();
141        Some(TimeRange::new(
142            DateTime::from_timestamp_nanos(slice[0]),
143            DateTime::from_timestamp_nanos(slice[slice.len() - 1]),
144        ))
145    }
146
147    pub fn append(&mut self, row: &Measure) -> Result<()> {
148        self.process_ids
149            .append(format!("{}", row.process.process_id))?;
150        self.stream_ids.append(&*row.stream_id)?;
151        self.block_ids.append(&*row.block_id)?;
152        self.insert_times.append_value(row.insert_time);
153        self.exes.append(&row.process.exe)?;
154        self.usernames.append(&row.process.username)?;
155        self.computers.append(&row.process.computer)?;
156        self.times.append_value(row.time);
157        self.targets.append(row.target)?;
158        self.names.append(row.name)?;
159        self.units.append(row.unit)?;
160        self.values.append_value(row.value);
161        self.properties.append_property_set(&row.properties)?;
162        self.process_properties.append(&*row.process.properties)?;
163        Ok(())
164    }
165
166    /// Append only per-entry variable data (optimized for batch processing)
167    pub fn append_entry_only(&mut self, row: &Measure) -> Result<()> {
168        // Only append fields that truly vary per metrics entry
169        self.times.append_value(row.time);
170        self.targets.append(row.target)?;
171        self.names.append(row.name)?;
172        self.units.append(row.unit)?;
173        self.values.append_value(row.value);
174        self.properties.append_property_set(&row.properties)?;
175        Ok(())
176    }
177
178    /// Batch fill all constant columns for all entries in block
179    pub fn fill_constant_columns(
180        &mut self,
181        process: &ProcessMetadata,
182        stream_id: &str,
183        block_id: &str,
184        insert_time: i64,
185        entry_count: usize,
186    ) -> Result<()> {
187        let process_id_str = format!("{}", process.process_id);
188
189        // For PrimitiveBuilder (insert_times): use append_slice for better performance
190        let insert_times_slice = vec![insert_time; entry_count];
191        self.insert_times.append_slice(&insert_times_slice);
192
193        // For BinaryDictionaryBuilder (process_properties): use append_n for same value
194        self.process_properties
195            .append_n(&**process.properties, entry_count)?;
196
197        // For StringDictionaryBuilder: use append_n for same values (optimal for constant data)
198        self.process_ids.append_n(&process_id_str, entry_count)?;
199        self.stream_ids.append_n(stream_id, entry_count)?;
200        self.block_ids.append_n(block_id, entry_count)?;
201        self.exes.append_n(&process.exe, entry_count)?;
202        self.usernames.append_n(&process.username, entry_count)?;
203        self.computers.append_n(&process.computer, entry_count)?;
204
205        Ok(())
206    }
207
208    pub fn finish(mut self) -> Result<RecordBatch> {
209        RecordBatch::try_new(
210            Arc::new(metrics_table_schema()),
211            vec![
212                Arc::new(self.process_ids.finish()),
213                Arc::new(self.stream_ids.finish()),
214                Arc::new(self.block_ids.finish()),
215                Arc::new(self.insert_times.finish().with_timezone_utc()),
216                Arc::new(self.exes.finish()),
217                Arc::new(self.usernames.finish()),
218                Arc::new(self.computers.finish()),
219                Arc::new(self.times.finish().with_timezone_utc()),
220                Arc::new(self.targets.finish()),
221                Arc::new(self.names.finish()),
222                Arc::new(self.units.finish()),
223                Arc::new(self.values.finish()),
224                Arc::new(self.properties.finish()?),
225                Arc::new(self.process_properties.finish()),
226            ],
227        )
228        .with_context(|| "building record batch")
229    }
230}
micromegas_analytics/metrics_table.rs

micromegas_analytics/
metrics_table.rs
