use std::collections::HashMap;
use std::sync::Arc;
use arrow_array::builder::{
ArrayBuilder, BinaryBuilder, BooleanBuilder, Float64Builder, Int64Builder,
StringDictionaryBuilder, UInt8Builder, UInt16Builder, UInt32Builder, UInt64Builder,
};
use arrow_array::types::{UInt16Type, UInt32Type};
use arrow_array::{Array, ArrayRef, RecordBatch};
use prost::Message;
use mira_proto::common::v1::{AnyValue, InstrumentationScope, KeyValue, any_value::Value};
use mira_proto::resource::v1::Resource;
use crate::error::{Error, Result};
use crate::identity::resource_key;
use crate::schema::{ATTRS, AttrType, DICT_CAP, RESOURCES};
pub struct DictColumn {
label: &'static str,
b: StringDictionaryBuilder<UInt16Type>,
n: usize,
}
impl DictColumn {
pub fn new(label: &'static str) -> Self {
Self {
label,
b: StringDictionaryBuilder::new(),
n: 0,
}
}
pub fn has_headroom(&self, n: usize) -> bool {
self.n + n <= DICT_CAP
}
pub fn append(&mut self, v: &str) -> Result<()> {
if v.is_empty() {
self.b.append_null();
return Ok(());
}
let k = self
.b
.append(v)
.map_err(|_| Error::DictionaryFull(self.label))?;
self.n = self.n.max(k as usize + 1);
Ok(())
}
pub fn finish(&self) -> ArrayRef {
Arc::new(self.b.finish_cloned())
}
}
pub struct AttrsBuilder {
label: &'static str,
parent_id: UInt32Builder,
key: StringDictionaryBuilder<UInt16Type>,
n_keys: usize,
ty: UInt8Builder,
str_: StringDictionaryBuilder<UInt32Type>,
n_str: usize,
str_bytes: usize,
int: Int64Builder,
double: Float64Builder,
bool_: BooleanBuilder,
bytes: BinaryBuilder,
ser: BinaryBuilder,
}
impl AttrsBuilder {
pub fn new(label: &'static str) -> Self {
Self {
label,
parent_id: UInt32Builder::new(),
key: StringDictionaryBuilder::new(),
n_keys: 0,
ty: UInt8Builder::new(),
str_: StringDictionaryBuilder::new(),
n_str: 0,
str_bytes: 0,
int: Int64Builder::new(),
double: Float64Builder::new(),
bool_: BooleanBuilder::new(),
bytes: BinaryBuilder::new(),
ser: BinaryBuilder::new(),
}
}
pub fn len(&self) -> usize {
self.ty.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn has_headroom(&self, n: usize) -> bool {
self.n_keys + n <= DICT_CAP
}
pub fn heap_bytes(&self) -> usize {
self.str_bytes + self.bytes.values_slice().len() + self.ser.values_slice().len()
}
pub fn append_all(&mut self, parent_id: u32, kvs: &[KeyValue]) -> Result<()> {
for kv in kvs {
self.append(parent_id, &kv.key, kv.value.as_ref())?;
}
Ok(())
}
pub fn append(&mut self, parent_id: u32, key: &str, value: Option<&AnyValue>) -> Result<()> {
let k = self
.key
.append(key)
.map_err(|_| Error::DictionaryFull(self.label))?;
self.n_keys = self.n_keys.max(k as usize + 1);
self.parent_id.append_value(parent_id);
let mut set = [false; 6];
let ty = match value.and_then(|v| v.value.as_ref()) {
None => AttrType::Empty,
Some(Value::StringValue(s)) => {
let k = self
.str_
.append(s)
.map_err(|_| Error::DictionaryFull(self.label))?;
if k as usize >= self.n_str {
self.n_str = k as usize + 1;
self.str_bytes += s.len();
}
set[0] = true;
AttrType::Str
}
Some(Value::IntValue(i)) => {
self.int.append_value(*i);
set[1] = true;
AttrType::Int
}
Some(Value::DoubleValue(d)) => {
self.double.append_value(*d);
set[2] = true;
AttrType::Double
}
Some(Value::BoolValue(b)) => {
self.bool_.append_value(*b);
set[3] = true;
AttrType::Bool
}
Some(Value::BytesValue(b)) => {
self.bytes.append_value(b);
set[4] = true;
AttrType::Bytes
}
Some(v @ Value::ArrayValue(_)) | Some(v @ Value::KvlistValue(_)) => {
let owned = AnyValue {
value: Some(v.clone()),
};
self.ser.append_value(owned.encode_to_vec());
set[5] = true;
if matches!(v, Value::ArrayValue(_)) {
AttrType::Slice
} else {
AttrType::Map
}
}
};
self.ty.append_value(ty as u8);
if !set[0] {
self.str_.append_null();
}
if !set[1] {
self.int.append_null();
}
if !set[2] {
self.double.append_null();
}
if !set[3] {
self.bool_.append_null();
}
if !set[4] {
self.bytes.append_null();
}
if !set[5] {
self.ser.append_null();
}
Ok(())
}
pub fn finish(&self) -> Result<RecordBatch> {
let cols: Vec<ArrayRef> = vec![
Arc::new(self.parent_id.finish_cloned()),
Arc::new(self.key.finish_cloned()),
Arc::new(self.ty.finish_cloned()),
Arc::new(self.str_.finish_cloned()),
Arc::new(self.int.finish_cloned()),
Arc::new(self.double.finish_cloned()),
Arc::new(self.bool_.finish_cloned()),
Arc::new(self.bytes.finish_cloned()),
Arc::new(self.ser.finish_cloned()),
];
Ok(RecordBatch::try_new(ATTRS.clone(), cols)?)
}
}
pub struct StrColumn<'a> {
keys: &'a arrow_array::UInt32Array,
values: &'a arrow_array::StringArray,
}
impl StrColumn<'_> {
pub fn value(&self, row: usize) -> &str {
self.values.value(self.keys.value(row) as usize)
}
pub fn is_valid(&self, row: usize) -> bool {
self.keys.is_valid(row)
}
}
pub fn str_column(b: &RecordBatch) -> StrColumn<'_> {
str_values(b.column(3))
}
pub fn str_values(col: &dyn Array) -> StrColumn<'_> {
use arrow_array::cast::AsArray;
let d = col.as_dictionary::<UInt32Type>();
StrColumn {
keys: d.keys(),
values: d.values().as_string::<i32>(),
}
}
pub fn index(tables: &[(&'static str, RecordBatch)]) -> Option<Vec<u8>> {
let mut keys = crate::bloom::Keys::default();
for (_, b) in tables {
if Arc::ptr_eq(&b.schema(), &ATTRS) {
index_table(&mut keys, b);
}
}
keys.build()
}
fn index_table(keys: &mut crate::bloom::Keys, b: &RecordBatch) {
use arrow_array::cast::AsArray;
use arrow_array::types::{Int64Type, UInt8Type};
let dict = b.column(1).as_dictionary::<UInt16Type>();
let names = dict.values().as_string::<i32>();
let codes = dict.keys().values();
let types = b.column(2).as_primitive::<UInt8Type>().values();
let strs = str_column(b);
let ints = b.column(4).as_primitive::<Int64Type>();
let bools = b.column(6).as_boolean();
const STR: u8 = AttrType::Str as u8;
const INT: u8 = AttrType::Int as u8;
const DOUBLE: u8 = AttrType::Double as u8;
const BOOL: u8 = AttrType::Bool as u8;
let mut buf = String::new();
for row in 0..b.num_rows() {
let name = names.value(codes[row] as usize);
let text: &str = match types[row] {
STR => strs.value(row),
INT => {
buf.clear();
use std::fmt::Write;
let _ = write!(buf, "{}", ints.value(row));
&buf
}
BOOL => {
if bools.value(row) {
"true"
} else {
"false"
}
}
DOUBLE => {
keys.flag(crate::bloom::HAS_DOUBLE);
continue;
}
_ => continue,
};
keys.insert(crate::bloom::attr_hash(name, text.as_bytes()));
}
}
pub struct ResourceScope {
resources: HashMap<Vec<u8>, u16>,
scopes: HashMap<Vec<u8>, u16>,
res_id: UInt16Builder,
res_key: UInt64Builder,
res_dropped: UInt32Builder,
pub resource_attrs: AttrsBuilder,
pub scope_attrs: AttrsBuilder,
}
impl Default for ResourceScope {
fn default() -> Self {
Self::new()
}
}
impl ResourceScope {
pub fn new() -> Self {
Self {
resources: HashMap::new(),
scopes: HashMap::new(),
res_id: UInt16Builder::new(),
res_key: UInt64Builder::new(),
res_dropped: UInt32Builder::new(),
resource_attrs: AttrsBuilder::new("resource_attrs.key"),
scope_attrs: AttrsBuilder::new("scope_attrs.key"),
}
}
pub fn has_headroom(
&self,
resources: usize,
scopes: usize,
res_kv: usize,
scope_kv: usize,
) -> bool {
self.resources.len() + resources <= DICT_CAP
&& self.scopes.len() + scopes <= DICT_CAP
&& self.resource_attrs.has_headroom(res_kv)
&& self.scope_attrs.has_headroom(scope_kv)
}
pub fn len(&self) -> usize {
self.resource_attrs.len() + self.scope_attrs.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn heap_bytes(&self) -> usize {
self.resource_attrs.heap_bytes() + self.scope_attrs.heap_bytes()
}
pub fn resource(&mut self, res: Option<&Resource>) -> Result<u16> {
let key = res.map(|r| r.encode_to_vec()).unwrap_or_default();
if let Some(&id) = self.resources.get(&key) {
return Ok(id);
}
let id = u16::try_from(self.resources.len())
.map_err(|_| Error::DictionaryFull("resource_id"))?;
self.resources.insert(key, id);
let attrs = res.map(|r| r.attributes.as_slice()).unwrap_or_default();
self.res_id.append_value(id);
self.res_key.append_value(resource_key(attrs));
self.res_dropped
.append_value(res.map(|r| r.dropped_attributes_count).unwrap_or(0));
self.resource_attrs.append_all(id as u32, attrs)?;
Ok(id)
}
pub fn scope(&mut self, scope: Option<&InstrumentationScope>) -> Result<u16> {
let key = scope.map(|s| s.encode_to_vec()).unwrap_or_default();
if let Some(&id) = self.scopes.get(&key) {
return Ok(id);
}
let id = u16::try_from(self.scopes.len()).map_err(|_| Error::DictionaryFull("scope_id"))?;
self.scopes.insert(key, id);
if let Some(s) = scope {
self.scope_attrs.append_all(id as u32, &s.attributes)?;
for (k, v) in [
("otel.scope.name", &s.name),
("otel.scope.version", &s.version),
] {
if !v.is_empty() {
let value = AnyValue {
value: Some(Value::StringValue(v.clone())),
};
self.scope_attrs.append(id as u32, k, Some(&value))?;
}
}
}
Ok(id)
}
pub fn finish(&self) -> Result<[(&'static str, RecordBatch); 3]> {
let cols: Vec<ArrayRef> = vec![
Arc::new(self.res_id.finish_cloned()),
Arc::new(self.res_key.finish_cloned()),
Arc::new(self.res_dropped.finish_cloned()),
];
let resources = RecordBatch::try_new(RESOURCES.clone(), cols)?;
Ok([
("resources", resources),
("resource_attrs", self.resource_attrs.finish()?),
("scope_attrs", self.scope_attrs.finish()?),
])
}
}
pub fn scope_kv(scope: Option<&InstrumentationScope>) -> usize {
scope.map_or(0, |s| s.attributes.len()) + 2
}
pub fn resource_kv(res: Option<&Resource>) -> usize {
res.map_or(0, |r| r.attributes.len())
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::Array;
use arrow_array::cast::AsArray;
use arrow_array::types::UInt8Type;
use mira_proto::common::v1::{ArrayValue, KeyValueList};
fn any(v: Value) -> AnyValue {
AnyValue { value: Some(v) }
}
#[test]
fn every_any_value_variant_fills_exactly_the_column_its_type_names() {
let mut b = AttrsBuilder::new("log_attrs.key");
assert!(b.is_empty(), "a fresh table has no rows");
let nested = KeyValueList {
values: vec![KeyValue {
key: "inner".into(),
value: Some(any(Value::IntValue(1))),
}],
};
let rows: [(&str, Option<AnyValue>, AttrType, Option<usize>); 9] = [
("absent", None, AttrType::Empty, None),
(
"unset",
Some(AnyValue { value: None }),
AttrType::Empty,
None,
),
(
"str",
Some(any(Value::StringValue("s".into()))),
AttrType::Str,
Some(3),
),
(
"int",
Some(any(Value::IntValue(-7))),
AttrType::Int,
Some(4),
),
(
"double",
Some(any(Value::DoubleValue(0.5))),
AttrType::Double,
Some(5),
),
(
"bool",
Some(any(Value::BoolValue(true))),
AttrType::Bool,
Some(6),
),
(
"bytes",
Some(any(Value::BytesValue(vec![0xde, 0xad].into()))),
AttrType::Bytes,
Some(7),
),
(
"slice",
Some(any(Value::ArrayValue(ArrayValue {
values: vec![any(Value::IntValue(1)), any(Value::StringValue("x".into()))],
}))),
AttrType::Slice,
Some(8),
),
(
"map",
Some(any(Value::KvlistValue(nested.clone()))),
AttrType::Map,
Some(8),
),
];
for (i, (key, value, _, _)) in rows.iter().enumerate() {
b.append(i as u32, key, value.as_ref()).expect("append");
}
assert_eq!(b.len(), rows.len());
assert!(!b.is_empty());
let batch = b.finish().expect("finish");
assert_eq!(batch.num_rows(), rows.len());
let types = batch.column(2).as_primitive::<UInt8Type>();
for (row, (key, _, ty, col)) in rows.iter().enumerate() {
assert_eq!(types.value(row), *ty as u8, "{key} stored the wrong type");
for c in 3..9 {
assert_eq!(
batch.column(c).is_null(row),
Some(c) != *col,
"{key}: column {c} nullness"
);
}
}
let ser = batch.column(8).as_binary::<i32>();
assert_eq!(
AnyValue::decode(ser.value(8)).expect("ser decodes"),
any(Value::KvlistValue(nested))
);
assert!(
b.heap_bytes() > ser.value(8).len(),
"the seal estimate must see every variable-width heap"
);
}
#[test]
fn a_repeated_attribute_value_is_stored_once() {
let prompt = "summarise the incident in one paragraph".repeat(64);
let mut b = AttrsBuilder::new("log_attrs.key");
for i in 0..1_000 {
b.append(
i,
"gen_ai.prompt",
Some(&any(Value::StringValue(prompt.clone()))),
)
.unwrap();
}
assert_eq!(
b.heap_bytes(),
prompt.len(),
"a thousand copies of one value are one value"
);
let batch = b.finish().unwrap();
let strs = str_column(&batch);
assert_eq!(strs.value(0), prompt);
assert_eq!(strs.value(999), prompt);
assert!(strs.is_valid(999));
assert!(
matches!(
batch.column(3).data_type(),
arrow_schema::DataType::Dictionary(k, _) if **k == arrow_schema::DataType::UInt32,
),
"the key width is the one the reader downcasts to"
);
b.append(
0,
"gen_ai.prompt",
Some(&any(Value::StringValue("no".into()))),
)
.unwrap();
assert_eq!(b.heap_bytes(), prompt.len() + 2);
}
#[test]
fn the_attribute_index_spells_every_comparable_value_the_way_a_query_will() {
let mut b = AttrsBuilder::new("log_attrs.key");
for (key, v) in [
("service.name", any(Value::StringValue("checkout".into()))),
("http.status", any(Value::IntValue(503))),
("canary", any(Value::BoolValue(true))),
("stable", any(Value::BoolValue(false))),
("ratio", any(Value::DoubleValue(0.25))),
("blob", any(Value::BytesValue(vec![1, 2].into()))),
] {
b.append(0, key, Some(&v)).expect("append");
}
b.append(0, "missing", None).expect("append");
let batch = b.finish().expect("finish");
let bytes = index(&[("log_attrs", batch)]).expect("an index over seven rows");
let f = crate::bloom::Filter::open(&bytes).expect("filter header");
for (key, text) in [
("service.name", "checkout"),
("http.status", "503"),
("canary", "true"),
("stable", "false"),
] {
assert!(
f.may_contain(crate::bloom::attr_hash(key, text.as_bytes())),
"{key}={text} was indexed as something else"
);
}
assert_eq!(f.flags & crate::bloom::HAS_DOUBLE, crate::bloom::HAS_DOUBLE);
let mut only_bytes = AttrsBuilder::new("log_attrs.key");
only_bytes
.append(0, "blob", Some(&any(Value::BytesValue(vec![9].into()))))
.expect("append");
assert!(index(&[("log_attrs", only_bytes.finish().expect("finish"))]).is_none());
}
#[test]
fn a_full_key_dictionary_is_an_error_that_leaves_the_block_sealable() {
let mut rs = ResourceScope::default();
assert!(rs.is_empty(), "a fresh preamble contributes no rows");
for i in 0..DICT_CAP {
rs.scope_attrs
.append(0, &format!("k{i}"), None)
.expect("headroom");
}
assert!(
!rs.has_headroom(1, 1, 0, 1),
"the hint must see the ceiling"
);
assert!(!rs.is_empty());
let scope = InstrumentationScope {
name: "payments".into(),
version: "1.2.3".into(),
..Default::default()
};
let e = rs.scope(Some(&scope)).expect_err("the dictionary is full");
assert!(matches!(e, Error::DictionaryFull("scope_attrs.key")), "{e}");
let tables = rs.finish().expect("a full table is still a sealable one");
assert_eq!(tables[2].0, "scope_attrs");
assert_eq!(tables[2].1.num_rows(), DICT_CAP);
}
#[test]
fn identical_resources_and_scopes_intern_to_one_row_and_different_ones_do_not() {
let mut rs = ResourceScope::new();
let res = |name: &str| Resource {
attributes: vec![KeyValue {
key: "service.name".into(),
value: Some(any(Value::StringValue(name.into()))),
}],
dropped_attributes_count: 0,
..Default::default()
};
assert_eq!(rs.resource(Some(&res("checkout"))).expect("resource"), 0);
assert_eq!(rs.resource(Some(&res("checkout"))).expect("resource"), 0);
assert_eq!(rs.resource(Some(&res("payments"))).expect("resource"), 1);
assert_eq!(rs.resource(None).expect("resource"), 2);
let scope = InstrumentationScope {
name: "tracer".into(),
..Default::default()
};
assert_eq!(rs.scope(Some(&scope)).expect("scope"), 0);
assert_eq!(rs.scope(Some(&scope)).expect("scope"), 0);
assert_eq!(rs.scope(None).expect("scope"), 1);
assert_eq!(rs.scope_attrs.len(), 1);
assert_eq!(
scope_kv(Some(&scope)),
2,
"the hint counts both, on purpose"
);
assert_eq!(resource_kv(Some(&res("checkout"))), 1);
assert_eq!(resource_kv(None), 0);
let tables = rs.finish().expect("finish");
let names: Vec<&str> = tables.iter().map(|(n, _)| *n).collect();
assert_eq!(names, ["resources", "resource_attrs", "scope_attrs"]);
assert_eq!(tables[0].1.num_rows(), 3, "three distinct resources");
assert_eq!(tables[1].1.num_rows(), 2, "and two of them carry one attr");
assert_eq!(rs.len(), 3);
assert!(rs.heap_bytes() > 0);
}
#[test]
fn a_dictionary_column_stores_the_unset_string_as_null_and_refuses_to_overflow() {
let mut d = DictColumn::new("logs.severity_text");
assert!(d.has_headroom(DICT_CAP));
d.append("").expect("empty");
d.append("ERROR").expect("value");
d.append("ERROR").expect("repeat");
let col = d.finish();
assert_eq!(col.len(), 3);
assert!(col.is_null(0), "proto3's unset string is not a slot");
assert!(d.has_headroom(DICT_CAP - 1));
assert!(
!d.has_headroom(DICT_CAP),
"one distinct value used, so one fewer fits"
);
}
}