diff --git a/crates/utils/re_mcap/src/layers/protobuf.rs b/crates/utils/re_mcap/src/layers/protobuf.rs index 84d66eb03233..bcf41cc31b3a 100644 --- a/crates/utils/re_mcap/src/layers/protobuf.rs +++ b/crates/utils/re_mcap/src/layers/protobuf.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use arrow::{ array::{ ArrayBuilder, BinaryBuilder, BooleanBuilder, FixedSizeListBuilder, Float32Builder, @@ -20,7 +18,7 @@ use crate::{Error, LayerIdentifier, MessageLayer}; struct ProtobufMessageParser { message_descriptor: MessageDescriptor, - fields: BTreeMap>>, + builder: FixedSizeListBuilder, } #[derive(Debug, thiserror::Error)] @@ -47,9 +45,6 @@ enum ProtobufError { #[error("unknown enum number {0}")] UnknownEnumNumber(i32), - #[error("unknown field name {0}")] - UnknownFieldName(String), - #[error("type {0} is not supported yet")] UnsupportedType(&'static str), @@ -59,19 +54,6 @@ enum ProtobufError { impl ProtobufMessageParser { fn new(num_rows: usize, message_descriptor: MessageDescriptor) -> Self { - let mut fields = BTreeMap::new(); - - // We recursively build up the Arrow builders for this particular message. - for field_descr in message_descriptor.fields() { - let name = field_descr.name().to_owned(); - let builder = arrow_builder_from_field(&field_descr); - fields.insert( - name, - FixedSizeListBuilder::with_capacity(builder, 1, num_rows), - ); - re_log::trace!("Added Arrow builder for fields: {}", field_descr.name()); - } - if message_descriptor.oneofs().len() > 0 { re_log::warn_once!( "`oneof` in schema {} is not supported yet.", @@ -84,9 +66,12 @@ impl ProtobufMessageParser { ); } + let struct_builder = struct_builder_from_message(&message_descriptor); + let builder = FixedSizeListBuilder::with_capacity(struct_builder, 1, num_rows); + Self { message_descriptor, - fields, + builder, } } } @@ -103,22 +88,31 @@ impl MessageParser for ProtobufMessageParser { }, )?; - // We always need to make sure to iterate over all our builders, adding null values whenever - // a field is missing from the message that we received. - for (field, builder) in &mut self.fields { - if let Some(val) = dynamic_message.get_field_by_name(field.as_str()) { + let struct_builder = self.builder.values(); + + for (ith_arrow_field, field_builder) in + struct_builder.field_builders_mut().iter_mut().enumerate() + { + // Protobuf fields are 1-indexed, so we need to map the i-th builder. + let protobuf_number = ith_arrow_field as u32 + 1; + + if let Some(val) = dynamic_message.get_field_by_number(protobuf_number) { let field = dynamic_message .descriptor() - .get_field_by_name(field) - .ok_or_else(|| ProtobufError::UnknownFieldName(field.to_owned()))?; - append_value(builder.values(), &field, val.as_ref())?; - builder.append(true); + .get_field(protobuf_number) + .ok_or_else(|| ProtobufError::MissingField { + field: protobuf_number, + })?; + append_value(field_builder, &field, val.as_ref())?; re_log::trace!("Field {}: Finished writing to builders", field.full_name()); } else { - builder.append(false); + append_null_to_builder(field_builder)?; } } + struct_builder.append(true); + self.builder.append(true); + Ok(()) } @@ -129,23 +123,19 @@ impl MessageParser for ProtobufMessageParser { let Self { message_descriptor, - fields, + mut builder, } = *self; let message_chunk = Chunk::from_auto_row_ids( ChunkId::new(), entity_path, timelines, - fields - .into_iter() - .map(|(field, mut builder)| { - ( - ComponentDescriptor::partial(field) - .with_builtin_archetype(message_descriptor.full_name()), - builder.finish().into(), - ) - }) - .collect(), + std::iter::once(( + ComponentDescriptor::partial("message") + .with_builtin_archetype(message_descriptor.full_name()), + builder.finish().into(), + )) + .collect(), ) .map_err(|err| Error::Other(anyhow::anyhow!(err)))?; @@ -166,6 +156,41 @@ fn downcast_err<'a, T: std::any::Any>( }) } +fn append_null_to_builder(builder: &mut dyn ArrayBuilder) -> Result<(), ProtobufError> { + // Try to append null by downcasting to known builder types + if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder.as_any_mut().downcast_mut::() { + b.append_null(); + } else if let Some(b) = builder + .as_any_mut() + .downcast_mut::>>() + { + b.append_null(); + } else { + return Err(ProtobufError::UnsupportedType( + "Unknown builder type for append_null", + )); + } + Ok(()) +} + fn append_value( builder: &mut dyn ArrayBuilder, field: &FieldDescriptor, diff --git a/crates/utils/re_mcap/src/layers/snapshots/re_mcap__layers__protobuf__test__two_simple_rows.snap b/crates/utils/re_mcap/src/layers/snapshots/re_mcap__layers__protobuf__test__two_simple_rows.snap index b766a3b27290..94f8e8f0ac0e 100644 --- a/crates/utils/re_mcap/src/layers/snapshots/re_mcap__layers__protobuf__test__two_simple_rows.snap +++ b/crates/utils/re_mcap/src/layers/snapshots/re_mcap__layers__protobuf__test__two_simple_rows.snap @@ -2,25 +2,25 @@ source: crates/utils/re_mcap/src/layers/protobuf.rs expression: "format!(\"{:-240}\", &chunks[0])" --- -┌────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ -│ METADATA: │ -│ * entity_path: /test_topic │ -│ * heap_size_bytes: [**REDACTED**] │ -│ * id: [**REDACTED**] │ -│ * version: [**REDACTED**] │ -├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ ┌─────────────────────────────────────┬─────────────────────────────┬─────────────────────────────┬───────────────────────────────────┬────────────────────────────────────┬─────────────────────────────────────┐ │ -│ │ RowId ┆ log_time ┆ publish_time ┆ com.example.Person:id ┆ com.example.Person:name ┆ com.example.Person:status │ │ -│ │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ -│ │ type: FixedSizeBinary[16] ┆ type: nullable Duration(ns) ┆ type: nullable Duration(ns) ┆ type: nullable List[nullable i32] ┆ type: nullable List[nullable Utf8] ┆ type: nullable List[nullable │ │ -│ │ ARROW:extension:metadata: ┆ index_name: log_time ┆ index_name: publish_time ┆ archetype: com.example.Person ┆ archetype: com.example.Person ┆ Struct[2]] │ │ -│ │ {"namespace":"row"} ┆ is_sorted: true ┆ is_sorted: true ┆ component: com.example.Person:id ┆ component: com.example.Person:name ┆ archetype: com.example.Person │ │ -│ │ ARROW:extension:name: TUID ┆ kind: index ┆ kind: index ┆ kind: data ┆ kind: data ┆ component: │ │ -│ │ is_sorted: true ┆ ┆ ┆ ┆ ┆ com.example.Person:status │ │ -│ │ kind: control ┆ ┆ ┆ ┆ ┆ kind: data │ │ -│ ╞═════════════════════════════════════╪═════════════════════════════╪═════════════════════════════╪═══════════════════════════════════╪════════════════════════════════════╪═════════════════════════════════════╡ │ -│ │ row_[**REDACTED**] ┆ PT0.000000042S ┆ PT0.000000042S ┆ [0] ┆ [Bob] ┆ [{name: INACTIVE, value: 2}] │ │ -│ ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ │ -│ │ row_[**REDACTED**] ┆ PT0.000000043S ┆ PT0.000000043S ┆ [123] ┆ [Alice] ┆ [{name: UNKNOWN, value: 0}] │ │ -│ └─────────────────────────────────────┴─────────────────────────────┴─────────────────────────────┴───────────────────────────────────┴────────────────────────────────────┴─────────────────────────────────────┘ │ -└────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ +┌──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ +│ METADATA: │ +│ * entity_path: /test_topic │ +│ * heap_size_bytes: [**REDACTED**] │ +│ * id: [**REDACTED**] │ +│ * version: [**REDACTED**] │ +├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ ┌───────────────────────────────────────────────┬─────────────────────────────┬─────────────────────────────┬──────────────────────────────────────────────────────────┐ │ +│ │ RowId ┆ log_time ┆ publish_time ┆ com.example.Person:message │ │ +│ │ --- ┆ --- ┆ --- ┆ --- │ │ +│ │ type: FixedSizeBinary[16] ┆ type: nullable Duration(ns) ┆ type: nullable Duration(ns) ┆ type: nullable List[nullable Struct[3]] │ │ +│ │ ARROW:extension:metadata: {"namespace":"row"} ┆ index_name: log_time ┆ index_name: publish_time ┆ archetype: com.example.Person │ │ +│ │ ARROW:extension:name: TUID ┆ is_sorted: true ┆ is_sorted: true ┆ component: com.example.Person:message │ │ +│ │ is_sorted: true ┆ kind: index ┆ kind: index ┆ kind: data │ │ +│ │ kind: control ┆ ┆ ┆ │ │ +│ ╞═══════════════════════════════════════════════╪═════════════════════════════╪═════════════════════════════╪══════════════════════════════════════════════════════════╡ │ +│ │ row_[**REDACTED**] ┆ PT0.000000042S ┆ PT0.000000042S ┆ [{name: Bob, id: 0, status: {name: INACTIVE, value: 2}}] │ │ +│ ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ │ +│ │ row_[**REDACTED**] ┆ PT0.000000043S ┆ PT0.000000043S ┆ [{name: Alice, id: 123, status: {name: UNKNOWN, value: │ │ +│ │ ┆ ┆ ┆ 0}}] │ │ +│ └───────────────────────────────────────────────┴─────────────────────────────┴─────────────────────────────┴──────────────────────────────────────────────────────────┘ │ +└──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘