1use crate::model::{self, Loaded};
22use anyhow::{Context, Result};
23use arrow::{
24 array::{ArrayRef, BooleanBuilder, Int64Builder, StringBuilder, StructArray, UInt64Builder},
25 buffer::NullBuffer,
26 datatypes::{DataType, Field, Fields, Schema},
27 record_batch::RecordBatch,
28};
29use backbeat::{
30 schema::{FieldRole, FieldType, Phase},
31 wire::{OwnedField, OwnedSchema},
32};
33use parquet::{
34 arrow::ArrowWriter,
35 file::{metadata::KeyValue, properties::WriterProperties},
36};
37use std::{
38 collections::{HashMap, HashSet},
39 fs::File,
40 path::Path,
41 sync::Arc,
42};
43
44pub fn to_parquet(dumps: &[Loaded], output: &Path, host: &str, zstd_level: i32) -> Result<usize> {
55 let mut schemas: Vec<OwnedSchema> = Vec::new();
58 let mut seen = HashSet::new();
59 for d in dumps {
60 for s in &d.schemas {
61 if seen.insert(s.id.get()) {
62 schemas.push(s.clone());
63 }
64 }
65 }
66 schemas.sort_by(|a, b| a.qualified_name.cmp(&b.qualified_name));
67 let by_id: HashMap<u64, usize> = schemas
68 .iter()
69 .enumerate()
70 .map(|(i, s)| (s.id.get(), i))
71 .collect();
72
73 let rows: Vec<Row> = model::unique_records(dumps)
78 .into_iter()
79 .map(|(d, r)| {
80 let id = d.schemas[r.schema_idx].id.get();
81 Row {
82 ts_nanos: r.ts_nanos,
83 instance_id: d.instance_id,
84 schema_idx: by_id[&id],
85 fields: &r.fields,
86 intern: &d.intern,
87 }
88 })
89 .collect();
90
91 let batch = build_batch(&schemas, &rows)?;
92
93 let footer_host = if !host.is_empty() {
96 host
97 } else {
98 dumps.first().map(|d| d.host.as_str()).unwrap_or("")
99 };
100
101 let mut seen_views = HashSet::new();
106 let mut tier2: Vec<String> = Vec::new();
107 for d in dumps {
108 for sql in &d.views {
109 if seen_views.insert(sql.clone()) {
110 tier2.push(sql.clone());
111 }
112 }
113 }
114 let ddl = crate::views::assemble(&schemas, &tier2);
115
116 write_parquet(output, &batch, footer_host, zstd_level, &ddl)?;
117
118 let sidecar = sidecar_path(output);
123 let parquet_name = output
124 .file_name()
125 .map(|n| n.to_string_lossy().into_owned())
126 .unwrap_or_else(|| output.to_string_lossy().into_owned());
127 let sidecar_body = format!("{}{ddl}", crate::views::bootstrap(&parquet_name));
128 std::fs::write(&sidecar, sidecar_body)
129 .with_context(|| format!("writing views sidecar {}", sidecar.display()))?;
130
131 Ok(rows.len())
132}
133
134fn sidecar_path(output: &Path) -> std::path::PathBuf {
136 output.with_extension("views.sql")
137}
138
139struct Row<'a> {
142 ts_nanos: u64,
143 instance_id: u64,
144 schema_idx: usize,
146 fields: &'a [u8],
147 intern: &'a HashMap<u32, String>,
148}
149
150enum Value {
152 U64(u64),
153 I64(i64),
154 Bool(bool),
155 Str(String),
156}
157
158fn decode_field(field: &OwnedField, bytes: &[u8], intern: &HashMap<u32, String>) -> Option<Value> {
163 let start = field.offset as usize;
164 let end = start + field.width as usize;
165 let slice = bytes.get(start..end)?;
166 if let Some(sentinel) = field.sentinel {
169 if field.width as usize <= 8 {
170 if let Some(raw) = read_uint(slice, field.width as usize) {
171 if raw == sentinel {
172 return None;
173 }
174 }
175 }
176 }
177 Some(match field.ty {
178 FieldType::U8 => Value::U64(slice[0] as u64),
179 FieldType::U16 => Value::U64(u16::from_le_bytes(slice.try_into().ok()?) as u64),
180 FieldType::U32 => Value::U64(u32::from_le_bytes(slice.try_into().ok()?) as u64),
181 FieldType::U64 => Value::U64(u64::from_le_bytes(slice.try_into().ok()?)),
182 FieldType::I8 => Value::I64(slice[0] as i8 as i64),
183 FieldType::I16 => Value::I64(i16::from_le_bytes(slice.try_into().ok()?) as i64),
184 FieldType::I32 => Value::I64(i32::from_le_bytes(slice.try_into().ok()?) as i64),
185 FieldType::I64 => Value::I64(i64::from_le_bytes(slice.try_into().ok()?)),
186 FieldType::Bool => Value::Bool(slice[0] != 0),
187 FieldType::Bytes => Value::Str(hex(slice)),
188 FieldType::Enum { repr } => {
189 let raw = read_uint(slice, repr as usize)?;
190 let label = field
191 .enum_labels
192 .iter()
193 .find(|l| l.value == raw)
194 .map(|l| l.label.clone())
195 .unwrap_or_else(|| raw.to_string());
196 Value::Str(label)
197 }
198 FieldType::Interned { .. } => {
199 let id = u32::from_le_bytes(slice.get(..4)?.try_into().ok()?);
200 Value::Str(intern.get(&id).cloned().unwrap_or_else(|| format!("#{id}")))
201 }
202 _ => Value::Str(hex(slice)),
204 })
205}
206
207fn read_uint(slice: &[u8], width: usize) -> Option<u64> {
209 let mut buf = [0u8; 8];
210 buf.get_mut(..width)?.copy_from_slice(slice.get(..width)?);
211 Some(u64::from_le_bytes(buf))
212}
213
214fn hex(bytes: &[u8]) -> String {
215 let mut s = String::with_capacity(bytes.len() * 2);
216 for b in bytes {
217 s.push_str(&format!("{b:02x}"));
218 }
219 s
220}
221
222fn arrow_type(ty: &FieldType) -> DataType {
224 match ty {
225 FieldType::U8 | FieldType::U16 | FieldType::U32 | FieldType::U64 => DataType::UInt64,
226 FieldType::I8 | FieldType::I16 | FieldType::I32 | FieldType::I64 => DataType::Int64,
227 FieldType::Bool => DataType::Boolean,
228 FieldType::Bytes | FieldType::Enum { .. } | FieldType::Interned { .. } => DataType::Utf8,
229 _ => DataType::Utf8,
231 }
232}
233
234fn display_names(schemas: &[OwnedSchema]) -> Vec<String> {
239 let mut counts: HashMap<&str, usize> = HashMap::new();
240 for s in schemas {
241 *counts.entry(s.qualified_name.as_str()).or_default() += 1;
242 }
243 schemas
244 .iter()
245 .map(|s| {
246 if counts[s.qualified_name.as_str()] > 1 {
247 format!("{}#{:016x}", s.qualified_name, s.id.get())
248 } else {
249 s.qualified_name.clone()
250 }
251 })
252 .collect()
253}
254
255fn is_promoted(role: FieldRole) -> bool {
258 matches!(
259 role,
260 FieldRole::Key | FieldRole::SpanId | FieldRole::ParentSpanId
261 )
262}
263
264fn field_metadata(f: &OwnedField) -> HashMap<String, String> {
267 let mut m = HashMap::new();
268 let role = match f.role {
269 FieldRole::None => None,
270 FieldRole::Key => Some("key"),
271 FieldRole::SpanId => Some("span_id"),
272 FieldRole::ParentSpanId => Some("parent_span_id"),
273 _ => None,
274 };
275 if let Some(role) = role {
276 m.insert("backbeat.role".to_string(), role.to_string());
277 }
278 if let Some(unit) = &f.unit {
279 m.insert("backbeat.unit".to_string(), unit.clone());
280 }
281 if let Some(sentinel) = f.sentinel {
282 m.insert("backbeat.sentinel".to_string(), sentinel.to_string());
285 }
286 if let Some(desc) = &f.description {
287 m.insert("backbeat.description".to_string(), desc.clone());
288 }
289 m
290}
291
292fn event_metadata(s: &OwnedSchema) -> HashMap<String, String> {
294 let mut m = HashMap::new();
295 let phase = match s.phase {
296 Phase::None => None,
297 Phase::Enter => Some("enter"),
298 Phase::Exit => Some("exit"),
299 _ => None,
300 };
301 if let Some(phase) = phase {
302 m.insert("backbeat.span".to_string(), phase.to_string());
303 }
304 if let Some(desc) = &s.description {
305 m.insert("backbeat.description".to_string(), desc.clone());
306 }
307 m
308}
309
310enum Col {
312 U64(UInt64Builder),
313 I64(Int64Builder),
314 Bool(BooleanBuilder),
315 Str(StringBuilder),
316}
317
318impl Col {
319 fn new(dt: &DataType) -> Self {
320 match dt {
321 DataType::UInt64 => Col::U64(UInt64Builder::new()),
322 DataType::Int64 => Col::I64(Int64Builder::new()),
323 DataType::Boolean => Col::Bool(BooleanBuilder::new()),
324 DataType::Utf8 => Col::Str(StringBuilder::new()),
325 other => unreachable!("unexpected column type {other:?}"),
326 }
327 }
328
329 fn append(&mut self, v: Value) {
330 match (self, v) {
331 (Col::U64(b), Value::U64(x)) => b.append_value(x),
332 (Col::I64(b), Value::I64(x)) => b.append_value(x),
333 (Col::Bool(b), Value::Bool(x)) => b.append_value(x),
334 (Col::Str(b), Value::Str(x)) => b.append_value(x),
335 (c, _) => c.append_null(),
337 }
338 }
339
340 fn append_null(&mut self) {
341 match self {
342 Col::U64(b) => b.append_null(),
343 Col::I64(b) => b.append_null(),
344 Col::Bool(b) => b.append_null(),
345 Col::Str(b) => b.append_null(),
346 }
347 }
348
349 fn finish(&mut self) -> ArrayRef {
350 match self {
351 Col::U64(b) => Arc::new(b.finish()),
352 Col::I64(b) => Arc::new(b.finish()),
353 Col::Bool(b) => Arc::new(b.finish()),
354 Col::Str(b) => Arc::new(b.finish()),
355 }
356 }
357}
358
359fn build_batch(schemas: &[OwnedSchema], rows: &[Row]) -> Result<RecordBatch> {
361 let names = display_names(schemas);
365
366 let mut key_names: Vec<String> = Vec::new();
371 let mut key_type: HashMap<String, DataType> = HashMap::new();
372 for s in schemas {
373 for f in s.fields.iter().filter(|f| is_promoted(f.role)) {
374 if !key_type.contains_key(&f.name) {
375 key_names.push(f.name.clone());
376 key_type.insert(f.name.clone(), arrow_type(&f.ty));
377 }
378 }
379 }
380
381 let mut seq = UInt64Builder::new();
383 let mut ts = UInt64Builder::new();
384 let mut event = StringBuilder::new();
385 let mut event_id = UInt64Builder::new();
386 let mut instance_id = UInt64Builder::new();
387 let mut key_cols: Vec<Col> = key_names.iter().map(|n| Col::new(&key_type[n])).collect();
388
389 struct EventCols {
391 children: Vec<(OwnedField, Col)>,
393 valid: Vec<bool>,
395 }
396 let mut event_cols: Vec<EventCols> = schemas
397 .iter()
398 .map(|s| EventCols {
399 children: s
400 .fields
401 .iter()
402 .filter(|f| !is_promoted(f.role))
403 .map(|f| (f.clone(), Col::new(&arrow_type(&f.ty))))
404 .collect(),
405 valid: Vec::with_capacity(rows.len()),
406 })
407 .collect();
408
409 for (i, row) in rows.iter().enumerate() {
411 let s = &schemas[row.schema_idx];
412 seq.append_value(i as u64);
413 ts.append_value(row.ts_nanos);
414 event.append_value(&names[row.schema_idx]);
415 event_id.append_value(s.id.get());
416 instance_id.append_value(row.instance_id);
417
418 for (name, col) in key_names.iter().zip(key_cols.iter_mut()) {
420 match s
421 .fields
422 .iter()
423 .find(|f| is_promoted(f.role) && &f.name == name)
424 {
425 Some(f) => match decode_field(f, row.fields, row.intern) {
426 Some(v) => col.append(v),
427 None => col.append_null(),
428 },
429 None => col.append_null(),
430 }
431 }
432
433 for (idx, ec) in event_cols.iter_mut().enumerate() {
435 let mine = idx == row.schema_idx;
436 ec.valid.push(mine);
437 for (f, col) in ec.children.iter_mut() {
438 if mine {
439 match decode_field(f, row.fields, row.intern) {
440 Some(v) => col.append(v),
441 None => col.append_null(),
442 }
443 } else {
444 col.append_null();
445 }
446 }
447 }
448 }
449
450 let mut fields: Vec<Field> = Vec::new();
452 let mut arrays: Vec<ArrayRef> = Vec::new();
453
454 fields.push(Field::new("seq", DataType::UInt64, false));
455 arrays.push(Arc::new(seq.finish()));
456 fields.push(Field::new("ts_nanos", DataType::UInt64, false));
457 arrays.push(Arc::new(ts.finish()));
458 fields.push(Field::new("event", DataType::Utf8, false));
459 arrays.push(Arc::new(event.finish()));
460 fields.push(Field::new("event_id", DataType::UInt64, false));
461 arrays.push(Arc::new(event_id.finish()));
462 fields.push(Field::new("instance_id", DataType::UInt64, false));
463 arrays.push(Arc::new(instance_id.finish()));
464
465 for (name, mut col) in key_names.iter().zip(key_cols) {
466 let decl = schemas
468 .iter()
469 .find_map(|s| s.fields.iter().find(|f| &f.name == name));
470 let mut field = Field::new(name, key_type[name].clone(), true);
471 if let Some(f) = decl {
472 field = field.with_metadata(field_metadata(f));
473 }
474 fields.push(field);
475 arrays.push(col.finish());
476 }
477
478 for (idx, mut ec) in event_cols.into_iter().enumerate() {
479 let s = &schemas[idx];
480 if ec.children.is_empty() {
485 continue;
486 }
487 let child_fields: Fields = ec
488 .children
489 .iter()
490 .map(|(f, _)| {
491 Field::new(&f.name, arrow_type(&f.ty), true).with_metadata(field_metadata(f))
492 })
493 .collect::<Vec<_>>()
494 .into();
495 let child_arrays: Vec<ArrayRef> = ec.children.iter_mut().map(|(_, c)| c.finish()).collect();
496 let nulls = NullBuffer::from(ec.valid);
497 let struct_array = StructArray::new(child_fields.clone(), child_arrays, Some(nulls));
498 let dt = DataType::Struct(child_fields);
499 let mut field = Field::new(&names[idx], dt, true);
502 field = field.with_metadata(event_metadata(s));
503 fields.push(field);
504 arrays.push(Arc::new(struct_array));
505 }
506
507 let schema = Arc::new(Schema::new(fields));
508 RecordBatch::try_new(schema, arrays).context("assembling record batch")
509}
510
511fn write_parquet(
519 output: &Path,
520 batch: &RecordBatch,
521 host: &str,
522 zstd_level: i32,
523 ddl: &str,
524) -> Result<()> {
525 let mut kv = vec![KeyValue::new(
526 "backbeat.format".to_string(),
527 "1".to_string(),
528 )];
529 if !host.is_empty() {
530 kv.push(KeyValue::new("backbeat.host".to_string(), host.to_string()));
531 }
532 if !ddl.is_empty() {
537 kv.push(KeyValue::new("backbeat.views".to_string(), ddl.to_string()));
538 }
539 let level = parquet::basic::ZstdLevel::try_new(zstd_level)
542 .with_context(|| format!("invalid zstd level {zstd_level} (valid range is 1–22)"))?;
543 let props = WriterProperties::builder()
544 .set_compression(parquet::basic::Compression::ZSTD(level))
545 .set_key_value_metadata(Some(kv))
546 .build();
547 let file = File::create(output).with_context(|| format!("creating {}", output.display()))?;
548 let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(props))?;
549 writer.write(batch)?;
550 writer.close()?;
551 Ok(())
552}