Skip to main content

micromegas_analytics/
log_entries_table.rs

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