1use std::collections::HashMap;
15use std::sync::Arc;
16
17use arrow_array::builder::{
18 ArrayBuilder, BinaryBuilder, BooleanBuilder, Float64Builder, Int64Builder,
19 StringDictionaryBuilder, UInt8Builder, UInt16Builder, UInt32Builder, UInt64Builder,
20};
21use arrow_array::types::{UInt16Type, UInt32Type};
22use arrow_array::{Array, ArrayRef, RecordBatch};
23use prost::Message;
24
25use mira_proto::common::v1::{AnyValue, InstrumentationScope, KeyValue, any_value::Value};
26use mira_proto::resource::v1::Resource;
27
28use crate::error::{Error, Result};
29use crate::identity::resource_key;
30use crate::schema::{ATTRS, AttrType, DICT_CAP, RESOURCES};
31
32pub struct DictColumn {
40 label: &'static str,
41 b: StringDictionaryBuilder<UInt16Type>,
42 n: usize,
43}
44
45impl DictColumn {
46 pub fn new(label: &'static str) -> Self {
47 Self {
48 label,
49 b: StringDictionaryBuilder::new(),
50 n: 0,
51 }
52 }
53
54 pub fn has_headroom(&self, n: usize) -> bool {
58 self.n + n <= DICT_CAP
59 }
60
61 pub fn append(&mut self, v: &str) -> Result<()> {
65 if v.is_empty() {
66 self.b.append_null();
67 return Ok(());
68 }
69 let k = self
70 .b
71 .append(v)
72 .map_err(|_| Error::DictionaryFull(self.label))?;
73 self.n = self.n.max(k as usize + 1);
74 Ok(())
75 }
76
77 pub fn finish(&self) -> ArrayRef {
86 Arc::new(self.b.finish_cloned())
87 }
88}
89
90pub struct AttrsBuilder {
93 label: &'static str,
96 parent_id: UInt32Builder,
97 key: StringDictionaryBuilder<UInt16Type>,
98 n_keys: usize,
102 ty: UInt8Builder,
103 str_: StringDictionaryBuilder<UInt32Type>,
107 n_str: usize,
113 str_bytes: usize,
114 int: Int64Builder,
115 double: Float64Builder,
116 bool_: BooleanBuilder,
117 bytes: BinaryBuilder,
118 ser: BinaryBuilder,
119}
120
121impl AttrsBuilder {
122 pub fn new(label: &'static str) -> Self {
123 Self {
124 label,
125 parent_id: UInt32Builder::new(),
126 key: StringDictionaryBuilder::new(),
127 n_keys: 0,
128 ty: UInt8Builder::new(),
129 str_: StringDictionaryBuilder::new(),
130 n_str: 0,
131 str_bytes: 0,
132 int: Int64Builder::new(),
133 double: Float64Builder::new(),
134 bool_: BooleanBuilder::new(),
135 bytes: BinaryBuilder::new(),
136 ser: BinaryBuilder::new(),
137 }
138 }
139
140 pub fn len(&self) -> usize {
141 self.ty.len()
142 }
143
144 pub fn is_empty(&self) -> bool {
145 self.len() == 0
146 }
147
148 pub fn has_headroom(&self, n: usize) -> bool {
152 self.n_keys + n <= DICT_CAP
153 }
154
155 pub fn heap_bytes(&self) -> usize {
159 self.str_bytes + self.bytes.values_slice().len() + self.ser.values_slice().len()
160 }
161
162 pub fn append_all(&mut self, parent_id: u32, kvs: &[KeyValue]) -> Result<()> {
164 for kv in kvs {
165 self.append(parent_id, &kv.key, kv.value.as_ref())?;
166 }
167 Ok(())
168 }
169
170 pub fn append(&mut self, parent_id: u32, key: &str, value: Option<&AnyValue>) -> Result<()> {
171 let k = self
177 .key
178 .append(key)
179 .map_err(|_| Error::DictionaryFull(self.label))?;
180 self.n_keys = self.n_keys.max(k as usize + 1);
181 self.parent_id.append_value(parent_id);
182
183 let mut set = [false; 6];
186 let ty = match value.and_then(|v| v.value.as_ref()) {
187 None => AttrType::Empty,
188 Some(Value::StringValue(s)) => {
189 let k = self
194 .str_
195 .append(s)
196 .map_err(|_| Error::DictionaryFull(self.label))?;
197 if k as usize >= self.n_str {
198 self.n_str = k as usize + 1;
199 self.str_bytes += s.len();
200 }
201 set[0] = true;
202 AttrType::Str
203 }
204 Some(Value::IntValue(i)) => {
205 self.int.append_value(*i);
206 set[1] = true;
207 AttrType::Int
208 }
209 Some(Value::DoubleValue(d)) => {
210 self.double.append_value(*d);
211 set[2] = true;
212 AttrType::Double
213 }
214 Some(Value::BoolValue(b)) => {
215 self.bool_.append_value(*b);
216 set[3] = true;
217 AttrType::Bool
218 }
219 Some(Value::BytesValue(b)) => {
220 self.bytes.append_value(b);
221 set[4] = true;
222 AttrType::Bytes
223 }
224 Some(v @ Value::ArrayValue(_)) | Some(v @ Value::KvlistValue(_)) => {
225 let owned = AnyValue {
226 value: Some(v.clone()),
227 };
228 self.ser.append_value(owned.encode_to_vec());
229 set[5] = true;
230 if matches!(v, Value::ArrayValue(_)) {
231 AttrType::Slice
232 } else {
233 AttrType::Map
234 }
235 }
236 };
237 self.ty.append_value(ty as u8);
238
239 if !set[0] {
240 self.str_.append_null();
241 }
242 if !set[1] {
243 self.int.append_null();
244 }
245 if !set[2] {
246 self.double.append_null();
247 }
248 if !set[3] {
249 self.bool_.append_null();
250 }
251 if !set[4] {
252 self.bytes.append_null();
253 }
254 if !set[5] {
255 self.ser.append_null();
256 }
257 Ok(())
258 }
259
260 pub fn finish(&self) -> Result<RecordBatch> {
262 let cols: Vec<ArrayRef> = vec![
263 Arc::new(self.parent_id.finish_cloned()),
264 Arc::new(self.key.finish_cloned()),
265 Arc::new(self.ty.finish_cloned()),
266 Arc::new(self.str_.finish_cloned()),
267 Arc::new(self.int.finish_cloned()),
268 Arc::new(self.double.finish_cloned()),
269 Arc::new(self.bool_.finish_cloned()),
270 Arc::new(self.bytes.finish_cloned()),
271 Arc::new(self.ser.finish_cloned()),
272 ];
273 Ok(RecordBatch::try_new(ATTRS.clone(), cols)?)
274 }
275}
276
277pub struct StrColumn<'a> {
289 keys: &'a arrow_array::UInt32Array,
290 values: &'a arrow_array::StringArray,
291}
292
293impl StrColumn<'_> {
294 pub fn value(&self, row: usize) -> &str {
295 self.values.value(self.keys.value(row) as usize)
296 }
297
298 pub fn is_valid(&self, row: usize) -> bool {
299 self.keys.is_valid(row)
300 }
301}
302
303pub fn str_column(b: &RecordBatch) -> StrColumn<'_> {
305 str_values(b.column(3))
306}
307
308pub fn str_values(col: &dyn Array) -> StrColumn<'_> {
311 use arrow_array::cast::AsArray;
312 let d = col.as_dictionary::<UInt32Type>();
313 StrColumn {
314 keys: d.keys(),
315 values: d.values().as_string::<i32>(),
316 }
317}
318
319pub fn index(tables: &[(&'static str, RecordBatch)]) -> Option<Vec<u8>> {
329 let mut keys = crate::bloom::Keys::default();
330 for (_, b) in tables {
331 if Arc::ptr_eq(&b.schema(), &ATTRS) {
332 index_table(&mut keys, b);
333 }
334 }
335 keys.build()
336}
337
338fn index_table(keys: &mut crate::bloom::Keys, b: &RecordBatch) {
339 use arrow_array::cast::AsArray;
340 use arrow_array::types::{Int64Type, UInt8Type};
341
342 let dict = b.column(1).as_dictionary::<UInt16Type>();
343 let names = dict.values().as_string::<i32>();
344 let codes = dict.keys().values();
345 let types = b.column(2).as_primitive::<UInt8Type>().values();
346 let strs = str_column(b);
347 let ints = b.column(4).as_primitive::<Int64Type>();
348 let bools = b.column(6).as_boolean();
349
350 const STR: u8 = AttrType::Str as u8;
351 const INT: u8 = AttrType::Int as u8;
352 const DOUBLE: u8 = AttrType::Double as u8;
353 const BOOL: u8 = AttrType::Bool as u8;
354
355 let mut buf = String::new();
358 for row in 0..b.num_rows() {
359 let name = names.value(codes[row] as usize);
360 let text: &str = match types[row] {
361 STR => strs.value(row),
362 INT => {
363 buf.clear();
364 use std::fmt::Write;
365 let _ = write!(buf, "{}", ints.value(row));
366 &buf
367 }
368 BOOL => {
369 if bools.value(row) {
370 "true"
371 } else {
372 "false"
373 }
374 }
375 DOUBLE => {
376 keys.flag(crate::bloom::HAS_DOUBLE);
377 continue;
378 }
379 _ => continue,
383 };
384 keys.insert(crate::bloom::attr_hash(name, text.as_bytes()));
385 }
386}
387
388pub struct ResourceScope {
401 resources: HashMap<Vec<u8>, u16>,
402 scopes: HashMap<Vec<u8>, u16>,
403 res_id: UInt16Builder,
404 res_key: UInt64Builder,
405 res_dropped: UInt32Builder,
406 pub resource_attrs: AttrsBuilder,
407 pub scope_attrs: AttrsBuilder,
408}
409
410impl Default for ResourceScope {
411 fn default() -> Self {
412 Self::new()
413 }
414}
415
416impl ResourceScope {
417 pub fn new() -> Self {
418 Self {
419 resources: HashMap::new(),
420 scopes: HashMap::new(),
421 res_id: UInt16Builder::new(),
422 res_key: UInt64Builder::new(),
423 res_dropped: UInt32Builder::new(),
424 resource_attrs: AttrsBuilder::new("resource_attrs.key"),
425 scope_attrs: AttrsBuilder::new("scope_attrs.key"),
426 }
427 }
428
429 pub fn has_headroom(
432 &self,
433 resources: usize,
434 scopes: usize,
435 res_kv: usize,
436 scope_kv: usize,
437 ) -> bool {
438 self.resources.len() + resources <= DICT_CAP
439 && self.scopes.len() + scopes <= DICT_CAP
440 && self.resource_attrs.has_headroom(res_kv)
441 && self.scope_attrs.has_headroom(scope_kv)
442 }
443
444 pub fn len(&self) -> usize {
446 self.resource_attrs.len() + self.scope_attrs.len()
447 }
448
449 pub fn is_empty(&self) -> bool {
450 self.len() == 0
451 }
452
453 pub fn heap_bytes(&self) -> usize {
454 self.resource_attrs.heap_bytes() + self.scope_attrs.heap_bytes()
455 }
456
457 pub fn resource(&mut self, res: Option<&Resource>) -> Result<u16> {
458 let key = res.map(|r| r.encode_to_vec()).unwrap_or_default();
459 if let Some(&id) = self.resources.get(&key) {
460 return Ok(id);
461 }
462 let id = u16::try_from(self.resources.len())
465 .map_err(|_| Error::DictionaryFull("resource_id"))?;
466 self.resources.insert(key, id);
467
468 let attrs = res.map(|r| r.attributes.as_slice()).unwrap_or_default();
469 self.res_id.append_value(id);
470 self.res_key.append_value(resource_key(attrs));
471 self.res_dropped
472 .append_value(res.map(|r| r.dropped_attributes_count).unwrap_or(0));
473 self.resource_attrs.append_all(id as u32, attrs)?;
474 Ok(id)
475 }
476
477 pub fn scope(&mut self, scope: Option<&InstrumentationScope>) -> Result<u16> {
478 let key = scope.map(|s| s.encode_to_vec()).unwrap_or_default();
479 if let Some(&id) = self.scopes.get(&key) {
480 return Ok(id);
481 }
482 let id = u16::try_from(self.scopes.len()).map_err(|_| Error::DictionaryFull("scope_id"))?;
483 self.scopes.insert(key, id);
484 if let Some(s) = scope {
485 self.scope_attrs.append_all(id as u32, &s.attributes)?;
486 for (k, v) in [
489 ("otel.scope.name", &s.name),
490 ("otel.scope.version", &s.version),
491 ] {
492 if !v.is_empty() {
493 let value = AnyValue {
494 value: Some(Value::StringValue(v.clone())),
495 };
496 self.scope_attrs.append(id as u32, k, Some(&value))?;
497 }
498 }
499 }
500 Ok(id)
501 }
502
503 pub fn finish(&self) -> Result<[(&'static str, RecordBatch); 3]> {
506 let cols: Vec<ArrayRef> = vec![
507 Arc::new(self.res_id.finish_cloned()),
508 Arc::new(self.res_key.finish_cloned()),
509 Arc::new(self.res_dropped.finish_cloned()),
510 ];
511 let resources = RecordBatch::try_new(RESOURCES.clone(), cols)?;
512 Ok([
513 ("resources", resources),
514 ("resource_attrs", self.resource_attrs.finish()?),
515 ("scope_attrs", self.scope_attrs.finish()?),
516 ])
517 }
518}
519
520pub fn scope_kv(scope: Option<&InstrumentationScope>) -> usize {
525 scope.map_or(0, |s| s.attributes.len()) + 2
526}
527
528pub fn resource_kv(res: Option<&Resource>) -> usize {
529 res.map_or(0, |r| r.attributes.len())
530}
531
532#[cfg(test)]
533mod tests {
534 use super::*;
535 use arrow_array::Array;
536 use arrow_array::cast::AsArray;
537 use arrow_array::types::UInt8Type;
538 use mira_proto::common::v1::{ArrayValue, KeyValueList};
539
540 fn any(v: Value) -> AnyValue {
541 AnyValue { value: Some(v) }
542 }
543
544 #[test]
553 fn every_any_value_variant_fills_exactly_the_column_its_type_names() {
554 let mut b = AttrsBuilder::new("log_attrs.key");
555 assert!(b.is_empty(), "a fresh table has no rows");
556
557 let nested = KeyValueList {
558 values: vec![KeyValue {
559 key: "inner".into(),
560 value: Some(any(Value::IntValue(1))),
561 }],
562 };
563 let rows: [(&str, Option<AnyValue>, AttrType, Option<usize>); 9] = [
569 ("absent", None, AttrType::Empty, None),
570 (
571 "unset",
572 Some(AnyValue { value: None }),
573 AttrType::Empty,
574 None,
575 ),
576 (
577 "str",
578 Some(any(Value::StringValue("s".into()))),
579 AttrType::Str,
580 Some(3),
581 ),
582 (
583 "int",
584 Some(any(Value::IntValue(-7))),
585 AttrType::Int,
586 Some(4),
587 ),
588 (
589 "double",
590 Some(any(Value::DoubleValue(0.5))),
591 AttrType::Double,
592 Some(5),
593 ),
594 (
595 "bool",
596 Some(any(Value::BoolValue(true))),
597 AttrType::Bool,
598 Some(6),
599 ),
600 (
601 "bytes",
602 Some(any(Value::BytesValue(vec![0xde, 0xad].into()))),
603 AttrType::Bytes,
604 Some(7),
605 ),
606 (
607 "slice",
608 Some(any(Value::ArrayValue(ArrayValue {
609 values: vec![any(Value::IntValue(1)), any(Value::StringValue("x".into()))],
610 }))),
611 AttrType::Slice,
612 Some(8),
613 ),
614 (
615 "map",
616 Some(any(Value::KvlistValue(nested.clone()))),
617 AttrType::Map,
618 Some(8),
619 ),
620 ];
621 for (i, (key, value, _, _)) in rows.iter().enumerate() {
622 b.append(i as u32, key, value.as_ref()).expect("append");
623 }
624 assert_eq!(b.len(), rows.len());
625 assert!(!b.is_empty());
626
627 let batch = b.finish().expect("finish");
628 assert_eq!(batch.num_rows(), rows.len());
629 let types = batch.column(2).as_primitive::<UInt8Type>();
630 for (row, (key, _, ty, col)) in rows.iter().enumerate() {
631 assert_eq!(types.value(row), *ty as u8, "{key} stored the wrong type");
632 for c in 3..9 {
633 assert_eq!(
634 batch.column(c).is_null(row),
635 Some(c) != *col,
636 "{key}: column {c} nullness"
637 );
638 }
639 }
640 let ser = batch.column(8).as_binary::<i32>();
643 assert_eq!(
644 AnyValue::decode(ser.value(8)).expect("ser decodes"),
645 any(Value::KvlistValue(nested))
646 );
647 assert!(
651 b.heap_bytes() > ser.value(8).len(),
652 "the seal estimate must see every variable-width heap"
653 );
654 }
655
656 #[test]
662 fn a_repeated_attribute_value_is_stored_once() {
663 let prompt = "summarise the incident in one paragraph".repeat(64);
664 let mut b = AttrsBuilder::new("log_attrs.key");
665 for i in 0..1_000 {
666 b.append(
667 i,
668 "gen_ai.prompt",
669 Some(&any(Value::StringValue(prompt.clone()))),
670 )
671 .unwrap();
672 }
673 assert_eq!(
674 b.heap_bytes(),
675 prompt.len(),
676 "a thousand copies of one value are one value"
677 );
678
679 let batch = b.finish().unwrap();
681 let strs = str_column(&batch);
682 assert_eq!(strs.value(0), prompt);
683 assert_eq!(strs.value(999), prompt);
684 assert!(strs.is_valid(999));
685 assert!(
686 matches!(
687 batch.column(3).data_type(),
688 arrow_schema::DataType::Dictionary(k, _) if **k == arrow_schema::DataType::UInt32,
689 ),
690 "the key width is the one the reader downcasts to"
691 );
692
693 b.append(
696 0,
697 "gen_ai.prompt",
698 Some(&any(Value::StringValue("no".into()))),
699 )
700 .unwrap();
701 assert_eq!(b.heap_bytes(), prompt.len() + 2);
702 }
703
704 #[test]
709 fn the_attribute_index_spells_every_comparable_value_the_way_a_query_will() {
710 let mut b = AttrsBuilder::new("log_attrs.key");
711 for (key, v) in [
712 ("service.name", any(Value::StringValue("checkout".into()))),
713 ("http.status", any(Value::IntValue(503))),
714 ("canary", any(Value::BoolValue(true))),
715 ("stable", any(Value::BoolValue(false))),
716 ("ratio", any(Value::DoubleValue(0.25))),
717 ("blob", any(Value::BytesValue(vec![1, 2].into()))),
718 ] {
719 b.append(0, key, Some(&v)).expect("append");
720 }
721 b.append(0, "missing", None).expect("append");
722 let batch = b.finish().expect("finish");
723
724 let bytes = index(&[("log_attrs", batch)]).expect("an index over seven rows");
725 let f = crate::bloom::Filter::open(&bytes).expect("filter header");
726 for (key, text) in [
727 ("service.name", "checkout"),
728 ("http.status", "503"),
729 ("canary", "true"),
730 ("stable", "false"),
731 ] {
732 assert!(
733 f.may_contain(crate::bloom::attr_hash(key, text.as_bytes())),
734 "{key}={text} was indexed as something else"
735 );
736 }
737 assert_eq!(f.flags & crate::bloom::HAS_DOUBLE, crate::bloom::HAS_DOUBLE);
740 let mut only_bytes = AttrsBuilder::new("log_attrs.key");
743 only_bytes
744 .append(0, "blob", Some(&any(Value::BytesValue(vec![9].into()))))
745 .expect("append");
746 assert!(index(&[("log_attrs", only_bytes.finish().expect("finish"))]).is_none());
747 }
748
749 #[test]
755 fn a_full_key_dictionary_is_an_error_that_leaves_the_block_sealable() {
756 let mut rs = ResourceScope::default();
757 assert!(rs.is_empty(), "a fresh preamble contributes no rows");
758
759 for i in 0..DICT_CAP {
760 rs.scope_attrs
761 .append(0, &format!("k{i}"), None)
762 .expect("headroom");
763 }
764 assert!(
765 !rs.has_headroom(1, 1, 0, 1),
766 "the hint must see the ceiling"
767 );
768 assert!(!rs.is_empty());
769
770 let scope = InstrumentationScope {
773 name: "payments".into(),
774 version: "1.2.3".into(),
775 ..Default::default()
776 };
777 let e = rs.scope(Some(&scope)).expect_err("the dictionary is full");
778 assert!(matches!(e, Error::DictionaryFull("scope_attrs.key")), "{e}");
779 let tables = rs.finish().expect("a full table is still a sealable one");
780 assert_eq!(tables[2].0, "scope_attrs");
781 assert_eq!(tables[2].1.num_rows(), DICT_CAP);
782 }
783
784 #[test]
790 fn identical_resources_and_scopes_intern_to_one_row_and_different_ones_do_not() {
791 let mut rs = ResourceScope::new();
792 let res = |name: &str| Resource {
793 attributes: vec![KeyValue {
794 key: "service.name".into(),
795 value: Some(any(Value::StringValue(name.into()))),
796 }],
797 dropped_attributes_count: 0,
798 ..Default::default()
799 };
800 assert_eq!(rs.resource(Some(&res("checkout"))).expect("resource"), 0);
801 assert_eq!(rs.resource(Some(&res("checkout"))).expect("resource"), 0);
802 assert_eq!(rs.resource(Some(&res("payments"))).expect("resource"), 1);
803 assert_eq!(rs.resource(None).expect("resource"), 2);
805
806 let scope = InstrumentationScope {
807 name: "tracer".into(),
808 ..Default::default()
809 };
810 assert_eq!(rs.scope(Some(&scope)).expect("scope"), 0);
811 assert_eq!(rs.scope(Some(&scope)).expect("scope"), 0);
812 assert_eq!(rs.scope(None).expect("scope"), 1);
813 assert_eq!(rs.scope_attrs.len(), 1);
816 assert_eq!(
817 scope_kv(Some(&scope)),
818 2,
819 "the hint counts both, on purpose"
820 );
821 assert_eq!(resource_kv(Some(&res("checkout"))), 1);
822 assert_eq!(resource_kv(None), 0);
823
824 let tables = rs.finish().expect("finish");
825 let names: Vec<&str> = tables.iter().map(|(n, _)| *n).collect();
826 assert_eq!(names, ["resources", "resource_attrs", "scope_attrs"]);
827 assert_eq!(tables[0].1.num_rows(), 3, "three distinct resources");
828 assert_eq!(tables[1].1.num_rows(), 2, "and two of them carry one attr");
829 assert_eq!(rs.len(), 3);
830 assert!(rs.heap_bytes() > 0);
831 }
832
833 #[test]
837 fn a_dictionary_column_stores_the_unset_string_as_null_and_refuses_to_overflow() {
838 let mut d = DictColumn::new("logs.severity_text");
839 assert!(d.has_headroom(DICT_CAP));
840 d.append("").expect("empty");
841 d.append("ERROR").expect("value");
842 d.append("ERROR").expect("repeat");
843 let col = d.finish();
844 assert_eq!(col.len(), 3);
845 assert!(col.is_null(0), "proto3's unset string is not a slot");
846 assert!(d.has_headroom(DICT_CAP - 1));
847 assert!(
848 !d.has_headroom(DICT_CAP),
849 "one distinct value used, so one fewer fits"
850 );
851 }
852}