Refactor the Arrow Mesh3D type to use zero-copy Buffers (#1691)

* Use arrow Buffer types for zero-copy mesh deserialization

* Use a Buffer for ColorRGBA as well

crates/re_log_types/src/component_types/mesh3d.rs

@@ -1,6 +1,5 @@
-use std::sync::Arc;
-
 use arrow2::array::{FixedSizeBinaryArray, MutableFixedSizeBinaryArray};
+use arrow2::buffer::Buffer;
 use arrow2::datatypes::DataType;
 use arrow2_convert::arrow_enable_vec_for_type;
 use arrow2_convert::deserialize::ArrowDeserialize;
@@ -9,7 +8,8 @@ use arrow2_convert::{serialize::ArrowSerialize, ArrowDeserialize, ArrowField, Ar
 
 use crate::msg_bundle::Component;
 
-use super::{ColorRGBA, FieldError, Vec4D};
+use super::arrow_convert_shims::BinaryBuffer;
+use super::{FieldError, Vec4D};
 
 // ----------------------------------------------------------------------------
 
@@ -137,31 +137,31 @@ pub enum RawMeshError {
 /// );
 /// ```
 #[derive(ArrowField, ArrowSerialize, ArrowDeserialize, Clone, Debug, PartialEq)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub struct RawMesh3D {
     pub mesh_id: MeshId,
 
     /// The flattened vertex positions array of this mesh.
     ///
     /// The length of this vector should always be divisible by three (since this is a 3D mesh).
     ///
-    /// If no indices are specified, then each triplet of vertex positions are intrpreted as a triangle
+    /// If no indices are specified, then each triplet of vertex positions are interpreted as a triangle
     /// and the length of this must be divisible by 9.
-    pub vertex_positions: Vec<f32>,
+    pub vertex_positions: Buffer<f32>,
 
     /// Per-vertex albedo colors.
-    pub vertex_colors: Option<Vec<ColorRGBA>>,
+    /// This is actually an encoded [`super::ColorRGBA`]
+    pub vertex_colors: Option<Buffer<u32>>,
 
     /// Optionally, the flattened normals array for this mesh.
     ///
     /// If specified, this must match the length of `Self::positions`.
-    pub vertex_normals: Option<Vec<f32>>,
+    pub vertex_normals: Option<Buffer<f32>>,
 
     /// Optionally, the flattened indices array for this mesh.
     ///
     /// Meshes are always triangle lists, i.e. the length of this vector should always be
     /// divisible by three.
-    pub indices: Option<Vec<u32>>,
+    pub indices: Option<Buffer<u32>>,
 
     /// Albedo factor applied to the final color of the mesh.
     ///
@@ -190,7 +190,7 @@ impl RawMesh3D {
                 return Err(RawMeshError::IndicesNotDivisibleBy3(indices.len()));
             }
 
-            for &index in indices {
+            for &index in indices.iter() {
                 if num_vertices <= index as usize {
                     return Err(RawMeshError::IndexOutOfBounds {
                         index,
@@ -257,13 +257,12 @@ impl RawMesh3D {
 /// );
 /// ```
 #[derive(Clone, Debug, PartialEq)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub struct EncodedMesh3D {
     pub mesh_id: MeshId,
 
     pub format: MeshFormat,
 
-    pub bytes: Arc<[u8]>,
+    pub bytes: BinaryBuffer,
 
     /// four columns of an affine transformation matrix
     pub transform: [[f32; 3]; 4],
@@ -276,7 +275,7 @@ pub struct EncodedMesh3DArrow {
 
     pub format: MeshFormat,
 
-    pub bytes: Vec<u8>,
+    pub bytes: BinaryBuffer,
 
     #[arrow_field(type = "arrow2_convert::field::FixedSizeVec<f32, 12>")]
     pub transform: Vec<f32>,
@@ -293,7 +292,7 @@ impl From<&EncodedMesh3D> for EncodedMesh3DArrow {
         Self {
             mesh_id: *mesh_id,
             format: *format,
-            bytes: bytes.as_ref().into(),
+            bytes: bytes.clone(),
             transform: transform.iter().flat_map(|c| c.iter().cloned()).collect(),
         }
     }
@@ -313,7 +312,7 @@ impl TryFrom<EncodedMesh3DArrow> for EncodedMesh3D {
         Ok(Self {
             mesh_id,
             format,
-            bytes: bytes.into(),
+            bytes,
             transform: [
                 transform.as_slice()[0..3].try_into()?,
                 transform.as_slice()[3..6].try_into()?,
@@ -406,7 +405,6 @@ impl std::fmt::Display for MeshFormat {
 /// ```
 #[derive(Clone, Debug, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)]
 #[arrow_field(type = "dense")]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub enum Mesh3D {
     Encoded(EncodedMesh3D),
     Raw(RawMesh3D),
@@ -436,14 +434,12 @@ mod tests {
     fn example_raw_mesh() -> RawMesh3D {
         let mesh = RawMesh3D {
             mesh_id: MeshId::random(),
-            vertex_positions: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 9.0, 10.0],
-            vertex_colors: Some(vec![
-                ColorRGBA(0xff0000ff),
-                ColorRGBA(0x00ff00ff),
-                ColorRGBA(0x0000ffff),
-            ]),
-            indices: vec![0, 1, 2].into(),
-            vertex_normals: vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 80.0, 90.0, 100.0].into(),
+            vertex_positions: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 9.0, 10.0].into(),
+            vertex_colors: Some(vec![0xff0000ff, 0x00ff00ff, 0x0000ffff].into()),
+            indices: Some(vec![0, 1, 2].into()),
+            vertex_normals: Some(
+                vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 80.0, 90.0, 100.0].into(),
+            ),
             albedo_factor: Vec4D([0.5, 0.5, 0.5, 1.0]).into(),
         };
         mesh.sanity_check().unwrap();
@@ -460,7 +456,7 @@ mod tests {
             let mesh_in = vec![Mesh3D::Encoded(EncodedMesh3D {
                 mesh_id: MeshId::random(),
                 format: MeshFormat::Glb,
-                bytes: std::sync::Arc::new([5, 9, 13, 95, 38, 42, 98, 17]),
+                bytes: vec![5, 9, 13, 95, 38, 42, 98, 17].into(),
                 transform: [
                     [1.0, 2.0, 3.0],
                     [4.0, 5.0, 6.0],

crates/re_viewer/src/misc/mesh_loader.rs

@@ -1,5 +1,4 @@
-use anyhow::anyhow;
-use re_log_types::{EncodedMesh3D, Mesh3D, MeshFormat, RawMesh3D};
+use re_log_types::{component_types::ColorRGBA, EncodedMesh3D, Mesh3D, MeshFormat, RawMesh3D};
 use re_renderer::{resource_managers::ResourceLifeTime, RenderContext, Rgba32Unmul};
 
 pub struct LoadedMesh {
@@ -71,7 +70,7 @@ impl LoadedMesh {
             transform,
         } = encoded_mesh;
 
-        let mut slf = Self::load_raw(name, *format, bytes, render_ctx)?;
+        let mut slf = Self::load_raw(name, *format, bytes.as_slice(), render_ctx)?;
 
         // TODO(cmc): Why are we creating the matrix twice here?
         let (scale, rotation, translation) =
@@ -106,8 +105,7 @@ impl LoadedMesh {
             albedo_factor,
         } = raw_mesh;
 
-        let vertex_positions: Vec<glam::Vec3> =
-            bytemuck::try_cast_vec(vertex_positions.clone()).map_err(|(err, _)| anyhow!(err))?;
+        let vertex_positions: &[glam::Vec3] = bytemuck::cast_slice(vertex_positions.as_slice());
         let num_positions = vertex_positions.len();
 
         let indices = if let Some(indices) = indices {
@@ -121,7 +119,7 @@ impl LoadedMesh {
         let vertex_colors = if let Some(vertex_colors) = vertex_colors {
             vertex_colors
                 .iter()
-                .map(|c| Rgba32Unmul::from_rgba_unmul_array(c.to_array()))
+                .map(|c| Rgba32Unmul::from_rgba_unmul_array(ColorRGBA(*c).to_array()))
                 .collect()
         } else {
             std::iter::repeat(Rgba32Unmul::WHITE)
@@ -148,8 +146,8 @@ impl LoadedMesh {
 
         let mesh = re_renderer::mesh::Mesh {
             label: name.clone().into(),
-            indices,
-            vertex_positions,
+            indices: indices.as_slice().into(),
+            vertex_positions: vertex_positions.into(),
             vertex_colors,
             vertex_normals,
             vertex_texcoords,

examples/rust/raw_mesh/src/main.rs

@@ -39,10 +39,10 @@ impl From<GltfPrimitive> for Mesh3D {
         let raw = RawMesh3D {
             mesh_id: MeshId::random(),
             albedo_factor: albedo_factor.map(Vec4D),
-            indices,
+            indices: indices.map(|i| i.into()),
             vertex_positions: vertex_positions.into_iter().flatten().collect(),
             vertex_normals: vertex_normals.map(|normals| normals.into_iter().flatten().collect()),
-            vertex_colors,
+            vertex_colors: vertex_colors.map(|colors| colors.into_iter().map(|c| c.0).collect()),
         };
 
         raw.sanity_check().unwrap();

rerun_py/src/python_bridge.rs

@@ -634,13 +634,13 @@ fn log_meshes(
                 [_, 3] => Some(
                     slice_from_np_array(&vertex_colors)
                         .chunks_exact(3)
-                        .map(|c| ColorRGBA::from_rgb(c[0], c[1], c[2]))
+                        .map(|c| ColorRGBA::from_rgb(c[0], c[1], c[2]).0)
                         .collect(),
                 ),
                 [_, 4] => Some(
                     slice_from_np_array(&vertex_colors)
                         .chunks_exact(4)
-                        .map(|c| ColorRGBA::from_unmultiplied_rgba(c[0], c[1], c[2], c[3]))
+                        .map(|c| ColorRGBA::from_unmultiplied_rgba(c[0], c[1], c[2], c[3]).0)
                         .collect(),
                 ),
                 shape => {
@@ -655,10 +655,10 @@ fn log_meshes(
 
         let raw = RawMesh3D {
             mesh_id: MeshId::random(),
-            vertex_positions: vertex_positions.as_array().to_vec(),
+            vertex_positions: vertex_positions.as_array().to_vec().into(),
             vertex_colors,
-            indices: indices.map(|indices| indices.as_array().to_vec()),
-            vertex_normals: normals.map(|normals| normals.as_array().to_vec()),
+            indices: indices.map(|indices| indices.as_array().to_vec().into()),
+            vertex_normals: normals.map(|normals| normals.as_array().to_vec().into()),
             albedo_factor,
         };
         raw.sanity_check()
@@ -708,7 +708,7 @@ fn log_mesh_file(
             )));
         }
     };
-    let bytes = bytes.into();
+    let bytes: Vec<u8> = bytes.into();
     let transform = if transform.is_empty() {
         [
             [1.0, 0.0, 0.0], // col 0
@@ -741,7 +741,7 @@ fn log_mesh_file(
     let mesh3d = Mesh3D::Encoded(EncodedMesh3D {
         mesh_id: MeshId::random(),
         format,
-        bytes,
+        bytes: bytes.into(),
         transform,
     });