Extract extra fields from ROS2 PointCloud2 messages

crates/utils/re_mcap/src/parsers/ros2msg/definitions/sensor_msgs.rs

@@ -110,11 +110,9 @@ pub struct Imu {
 }
 
 #[derive(Copy, Clone, Debug, Serialize, Deserialize)]
-#[serde(from = "u8", into = "u8")]
+#[serde(try_from = "u8", into = "u8")]
 #[repr(u8)]
 pub enum PointFieldDatatype {
-    /// Does not exist in original spec.
-    Unknown = 0,
     Int8 = 1,
     UInt8 = 2,
     Int16 = 3,
@@ -125,9 +123,15 @@ pub enum PointFieldDatatype {
     Float64 = 8,
 }
 
-impl From<u8> for PointFieldDatatype {
-    fn from(value: u8) -> Self {
-        match value {
+#[derive(Debug, thiserror::Error)]
+#[error("unknown point field datatype: {0}")]
+pub struct UnknownPointFieldDatatype(u8);
+
+impl TryFrom<u8> for PointFieldDatatype {
+    type Error = UnknownPointFieldDatatype;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        Ok(match value {
             1 => Self::Int8,
             2 => Self::UInt8,
             3 => Self::Int16,
@@ -136,8 +140,8 @@ impl From<u8> for PointFieldDatatype {
             6 => Self::UInt32,
             7 => Self::Float32,
             8 => Self::Float64,
-            _ => Self::Unknown,
-        }
+            other => Err(UnknownPointFieldDatatype(other))?,
+        })
     }
 }
 

crates/utils/re_mcap/src/parsers/ros2msg/sensor_msgs/point_cloud_2.rs

@@ -1,18 +1,20 @@
 use std::io::Cursor;
 
 use super::super::definitions::sensor_msgs::{self, PointField, PointFieldDatatype};
+use anyhow::Context as _;
 use arrow::{
     array::{
-        BooleanBuilder, FixedSizeListBuilder, ListBuilder, StringBuilder, StructBuilder,
-        UInt8Builder, UInt32Builder,
+        ArrayBuilder, BooleanBuilder, FixedSizeListBuilder, Float32Builder, Float64Builder,
+        Int8Builder, Int16Builder, Int32Builder, ListBuilder, StringBuilder, StructBuilder,
+        UInt8Builder, UInt16Builder, UInt32Builder,
     },
     datatypes::{DataType, Field, Fields},
 };
 use byteorder::{BigEndian, LittleEndian, ReadBytesExt as _};
 use re_chunk::{Chunk, ChunkComponents, ChunkId};
 use re_types::{
-    AsComponents as _, Component as _, ComponentDescriptor, SerializedComponentColumn, archetypes,
-    components, reflection::ComponentDescriptorExt as _,
+    Archetype as _, AsComponents as _, Component as _, ComponentDescriptor,
+    SerializedComponentColumn, archetypes, components, reflection::ComponentDescriptorExt as _,
 };
 use std::collections::HashMap;
 
@@ -38,6 +40,8 @@ pub struct PointCloud2MessageParser {
     data: FixedSizeListBuilder<ListBuilder<UInt8Builder>>,
     is_dense: FixedSizeListBuilder<BooleanBuilder>,
 
+    extracted_fields: Vec<(String, ListBuilder<Box<dyn ArrayBuilder>>)>,
+
     // We lazily create this, only if we can interpret the point cloud semantically.
     // For now, this is the case if there are fields with names `x`,`y`, and `z` present.
     points_3ds: Option<Vec<archetypes::Points3D>>,
@@ -80,15 +84,29 @@ impl Ros2MessageParser for PointCloud2MessageParser {
             data: blob_list_builder(num_rows),
             is_dense: fixed_size_list_builder(1, num_rows),
 
+            extracted_fields: Default::default(),
+
             points_3ds: None,
         }
     }
 }
 
+fn builder_from_datatype(datatype: PointFieldDatatype) -> Box<dyn ArrayBuilder> {
+    match datatype {
+        PointFieldDatatype::Int8 => Box::new(Int8Builder::new()),
+        PointFieldDatatype::UInt8 => Box::new(UInt8Builder::new()),
+        PointFieldDatatype::Int16 => Box::new(Int16Builder::new()),
+        PointFieldDatatype::UInt16 => Box::new(UInt16Builder::new()),
+        PointFieldDatatype::Int32 => Box::new(Int32Builder::new()),
+        PointFieldDatatype::UInt32 => Box::new(UInt32Builder::new()),
+        PointFieldDatatype::Float32 => Box::new(Float32Builder::new()),
+        PointFieldDatatype::Float64 => Box::new(Float64Builder::new()),
+    }
+}
+
 fn access(data: &[u8], datatype: PointFieldDatatype, is_big_endian: bool) -> std::io::Result<f32> {
     let mut rdr = Cursor::new(data);
     match (is_big_endian, datatype) {
-        (_, PointFieldDatatype::Unknown) => Ok(0f32), // Not in the original spec.
         (_, PointFieldDatatype::UInt8) => rdr.read_u8().map(|x| x as f32),
         (_, PointFieldDatatype::Int8) => rdr.read_i8().map(|x| x as f32),
         (true, PointFieldDatatype::Int16) => rdr.read_i16::<BigEndian>().map(|x| x as f32),
@@ -106,6 +124,21 @@ fn access(data: &[u8], datatype: PointFieldDatatype, is_big_endian: bool) -> std
     }
 }
 
+impl From<PointFieldDatatype> for DataType {
+    fn from(value: PointFieldDatatype) -> Self {
+        match value {
+            PointFieldDatatype::Int8 => Self::Int8,
+            PointFieldDatatype::UInt8 => Self::UInt8,
+            PointFieldDatatype::Int16 => Self::Int16,
+            PointFieldDatatype::UInt16 => Self::UInt16,
+            PointFieldDatatype::Int32 => Self::Int32,
+            PointFieldDatatype::UInt32 => Self::UInt32,
+            PointFieldDatatype::Float32 => Self::Float32,
+            PointFieldDatatype::Float64 => Self::Float64,
+        }
+    }
+}
+
 pub struct Position3DIter<'a> {
     point_iter: std::slice::ChunksExact<'a, u8>,
     is_big_endian: bool,
@@ -172,6 +205,127 @@ impl Iterator for Position3DIter<'_> {
     }
 }
 
+fn add_field_value(
+    builder: &mut Box<dyn ArrayBuilder>,
+    field: &PointField,
+    is_big_endian: bool,
+    data: &[u8],
+) -> anyhow::Result<()> {
+    let mut rdr = Cursor::new(data);
+    match field.datatype {
+        PointFieldDatatype::Int8 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<Int8Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `Int8Builder`", field.datatype)
+                })?;
+            let val = rdr.read_i8()?;
+            builder.append_value(val);
+        }
+        PointFieldDatatype::UInt8 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<UInt8Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `UInt8Builder`", field.datatype)
+                })?;
+            let val = rdr.read_u8()?;
+            builder.append_value(val);
+        }
+        PointFieldDatatype::Int16 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<Int16Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `Int16Builder`", field.datatype)
+                })?;
+            let val = if is_big_endian {
+                rdr.read_i16::<BigEndian>()?
+            } else {
+                rdr.read_i16::<LittleEndian>()?
+            };
+            builder.append_value(val);
+        }
+        PointFieldDatatype::UInt16 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<UInt16Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `UInt16Builder`", field.datatype)
+                })?;
+            let val = if is_big_endian {
+                rdr.read_u16::<BigEndian>()?
+            } else {
+                rdr.read_u16::<LittleEndian>()?
+            };
+            builder.append_value(val);
+        }
+
+        PointFieldDatatype::Int32 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<Int32Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `Int32Builder`", field.datatype)
+                })?;
+
+            let val = if is_big_endian {
+                rdr.read_i32::<BigEndian>()?
+            } else {
+                rdr.read_i32::<LittleEndian>()?
+            };
+            builder.append_value(val);
+        }
+        PointFieldDatatype::UInt32 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<UInt32Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `UInt16Builder`", field.datatype)
+                })?;
+            let val = if is_big_endian {
+                rdr.read_u32::<BigEndian>()?
+            } else {
+                rdr.read_u32::<LittleEndian>()?
+            };
+            builder.append_value(val);
+        }
+
+        PointFieldDatatype::Float32 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<Float32Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `Float32Builder`", field.datatype)
+                })?;
+            let val = if is_big_endian {
+                rdr.read_f32::<BigEndian>()?
+            } else {
+                rdr.read_f32::<LittleEndian>()?
+            };
+            builder.append_value(val);
+        }
+
+        PointFieldDatatype::Float64 => {
+            let builder = builder
+                .as_any_mut()
+                .downcast_mut::<Float64Builder>()
+                .with_context(|| {
+                    format!("found datatype {:?}, but `Float64Builder`", field.datatype)
+                })?;
+            let val = if is_big_endian {
+                rdr.read_f64::<BigEndian>()?
+            } else {
+                rdr.read_f64::<LittleEndian>()?
+            };
+            builder.append_value(val);
+        }
+    }
+
+    Ok(())
+}
+
 impl MessageParser for PointCloud2MessageParser {
     fn append(&mut self, ctx: &mut ParserContext, msg: &mcap::Message<'_>) -> anyhow::Result<()> {
         let point_cloud = cdr::try_decode_message::<sensor_msgs::PointCloud2>(msg.data.as_ref())
@@ -194,6 +348,8 @@ impl MessageParser for PointCloud2MessageParser {
             data,
             is_dense,
 
+            extracted_fields,
+
             points_3ds,
         } = self;
 
@@ -207,6 +363,34 @@ impl MessageParser for PointCloud2MessageParser {
             &point_cloud.fields,
         );
 
+        // We lazily initialize the builders that store the extracted fields from
+        // the blob when we receive the first message.
+        if extracted_fields.len() != point_cloud.fields.len() {
+            *extracted_fields = point_cloud
+                .fields
+                .iter()
+                .map(|field| {
+                    (
+                        field.name.clone(),
+                        ListBuilder::new(builder_from_datatype(field.datatype)),
+                    )
+                })
+                .collect();
+        }
+
+        for point in point_cloud.data.chunks(point_cloud.point_step as usize) {
+            for (field, (_name, builder)) in
+                point_cloud.fields.iter().zip(extracted_fields.iter_mut())
+            {
+                let field_builder = builder.values();
+                add_field_value(field_builder, field, point_cloud.is_bigendian, point)?;
+            }
+        }
+
+        for (_name, builder) in extracted_fields {
+            builder.append(true);
+        }
+
         if let Some(position_iter) = position_iter {
             points_3ds
                 .get_or_insert_with(|| Vec::with_capacity(*num_rows))
@@ -290,12 +474,14 @@ impl MessageParser for PointCloud2MessageParser {
             mut data,
             mut is_dense,
 
+            extracted_fields: points,
+
             points_3ds,
         } = *self;
 
         let mut chunks = Vec::new();
 
-        for (i, points_3d) in points_3ds.into_iter().enumerate() {
+        for (i, points_3d) in points_3ds.iter().enumerate() {
             let timelines = timelines
                 .iter()
                 .map(|(timeline, time_col)| (*timeline, time_col.row_sliced(i, 1).clone()))
@@ -365,6 +551,24 @@ impl MessageParser for PointCloud2MessageParser {
                 ),
             ]
             .into_iter()
+            .chain(points.into_iter().filter_map(|(name, mut builder)| {
+                // We only extract additional fields when we have a `Points3d`
+                // archetype to attach them to. In that case we're not interested
+                // in the other components.
+                // TODO(grtlr): It would be nice to never initialize the unnecessary builders
+                // in the first place. But, we'll soon move the semantic extraction of `Points3d`
+                // into a different layer anyways, making that optimization obsolete.
+                points_3ds.as_ref()?;
+                if ["x", "y", "z"].contains(&name.as_str()) {
+                    None
+                } else {
+                    Some((
+                        ComponentDescriptor::partial(name.clone())
+                            .with_builtin_archetype(archetypes::Points3D::name()),
+                        builder.finish(),
+                    ))
+                }
+            }))
             .collect(),
         )?;