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
17pub 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 Field::new(
92 "audience",
93 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
94 false,
95 ),
96 ])
97}
98
99pub 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 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 pub fn append_entry_only(&mut self, row: &Measure) -> Result<()> {
181 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 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 let insert_times_slice = vec![insert_time; entry_count];
204 self.insert_times.append_slice(&insert_times_slice);
205
206 self.process_properties
208 .append_n(&**process.properties, entry_count)?;
209
210 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}