End-to-end tagging: basic helpers for Python & Rust (#8332)

Introduce some very basic helpers so that overriding tags doesn't
require a PhD.

The goal here is not _great_, it's simply _not completely horrible_.

## Python

Before:
```py
class CustomPosition3DBatch(rr.ComponentBatchLike):
    def __init__(self: Any, positions: npt.ArrayLike) -> None:
        self.position = rr.components.Position3DBatch(positions)

    def component_descriptor(self) -> rr.ComponentDescriptor:
        return rr.ComponentDescriptor(
            "user.CustomPosition3D",
            archetype_name="user.CustomPoints3D",
            archetype_field_name="custom_positions",
        )

    def as_arrow_array(self) -> pa.Array:
        return self.position.as_arrow_array()


class CustomPoints3D(rr.AsComponents):
    def __init__(self: Any, positions: npt.ArrayLike, colors: npt.ArrayLike) -> None:
        self.positions = CustomPosition3DBatch(positions)
        self.colors = rr.components.ColorBatch(colors)

    def as_component_batches(self) -> Iterable[rr.ComponentBatchLike]:
        return (
            [rr.IndicatorComponentBatch("user.CustomPoints3D")]
            + [
                rr.DescribedComponentBatch(
                    self.positions,
                    self.positions.component_descriptor().or_with_overrides(
                        archetype_name="user.CustomPoints3D", archetype_field_name="custom_positions"
                    ),
                )
            ]
            + [
                rr.DescribedComponentBatch(
                    self.colors,
                    self.colors.component_descriptor().or_with_overrides(
                        archetype_name="user.CustomPoints3D", archetype_field_name="colors"
                    ),
                )
            ]
        )

```

After:
```py
class CustomPoints3D(rr.AsComponents):
    def __init__(self: Any, positions: npt.ArrayLike, colors: npt.ArrayLike) -> None:
        self.positions = rr.components.Position3DBatch(positions).with_descriptor(
            rr.ComponentDescriptor(
                "user.CustomPosition3D",
                archetype_name="user.CustomPoints3D",
                archetype_field_name="custom_positions",
            )
        )
        self.colors = rr.components.ColorBatch(colors).or_with_descriptor_overrides(
            archetype_name="user.CustomPoints3D",
            archetype_field_name="colors",
        )

    def as_component_batches(self) -> Iterable[rr.ComponentBatchLike]:
        return [rr.IndicatorComponentBatch("user.CustomPoints3D"), self.positions, self.colors]
```

## Rust

Before:
```rust
#[derive(Debug, Clone, Copy)]
struct CustomPosition3D(rerun::components::Position3D);

impl rerun::SizeBytes for CustomPosition3D {
    #[inline]
    fn heap_size_bytes(&self) -> u64 {
        0
    }
}

impl rerun::Loggable for CustomPosition3D {
    #[inline]
    fn arrow2_datatype() -> arrow2::datatypes::DataType {
        rerun::components::Position3D::arrow2_datatype()
    }

    #[inline]
    fn to_arrow2_opt<'a>(
        data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
    ) -> rerun::SerializationResult<Box<dyn arrow2::array::Array>>
    where
        Self: 'a,
    {
        rerun::components::Position3D::to_arrow2_opt(
            data.into_iter().map(|opt| opt.map(Into::into).map(|c| c.0)),
        )
    }
}

impl rerun::Component for CustomPosition3D {
    #[inline]
    fn descriptor() -> ComponentDescriptor {
        ComponentDescriptor::new("user.CustomPosition3D")
    }
}

struct CustomPoints3D {
    positions: Vec<CustomPosition3D>,
    colors: Option<Vec<rerun::components::Color>>,
}

impl CustomPoints3D {
    fn indicator() -> rerun::NamedIndicatorComponent {
        rerun::NamedIndicatorComponent("user.CustomPoints3DIndicator".into())
    }

    fn overridden_position_descriptor() -> ComponentDescriptor {
        CustomPosition3D::descriptor()
            .or_with_archetype_name(|| "user.CustomPoints3D".into())
            .or_with_archetype_field_name(|| "custom_positions".into())
    }

    fn overridden_color_descriptor() -> ComponentDescriptor {
        rerun::components::Color::descriptor()
            .or_with_archetype_name(|| "user.CustomPoints3D".into())
            .or_with_archetype_field_name(|| "colors".into())
    }
}

impl rerun::AsComponents for CustomPoints3D {
    fn as_component_batches(&self) -> Vec<rerun::MaybeOwnedComponentBatch<'_>> {
        [
            Some(Self::indicator().to_batch()),
            Some(
                rerun::MaybeOwnedComponentBatch::new(&self.positions as &dyn rerun::ComponentBatch)
                    .with_descriptor_override(Self::overridden_position_descriptor()),
            ),
            self.colors.as_ref().map(|colors| {
                rerun::MaybeOwnedComponentBatch::new(colors as &dyn rerun::ComponentBatch)
                    .with_descriptor_override(Self::overridden_color_descriptor())
            }),
        ]
        .into_iter()
        .flatten()
        .collect()
    }
}
```

After:
```rust
struct CustomPoints3D {
    positions: Vec<rerun::components::Position3D>,
    colors: Option<Vec<rerun::components::Color>>,
}

impl CustomPoints3D {
    fn indicator() -> rerun::NamedIndicatorComponent {
        rerun::NamedIndicatorComponent("user.CustomPoints3DIndicator".into())
    }

    fn overridden_position_descriptor() -> ComponentDescriptor {
        ComponentDescriptor {
            archetype_name: Some("user.CustomPoints3D".into()),
            archetype_field_name: Some("custom_positions".into()),
            component_name: "user.CustomPosition3D".into(),
        }
    }

    fn overridden_color_descriptor() -> ComponentDescriptor {
        <rerun::components::Color as rerun::Component>::descriptor()
            .or_with_archetype_name(|| "user.CustomPoints3D".into())
            .or_with_archetype_field_name(|| "colors".into())
    }
}

impl rerun::AsComponents for CustomPoints3D {
    fn as_component_batches(&self) -> Vec<rerun::MaybeOwnedComponentBatch<'_>> {
        [
            Some(Self::indicator().to_batch()),
            Some(
                self.positions
                    .with_descriptor(Self::overridden_position_descriptor()),
            ),
            self.colors
                .as_ref()
                .map(|colors| colors.with_descriptor(Self::overridden_color_descriptor())),
        ]
        .into_iter()
        .flatten()
        .collect()
    }
}
```

## C++

Too complicated, @Wumpf and I made the wise decision of giving up
:upside_down_face:

---------

Co-authored-by: Andreas Reich <[email protected]>

Author: teh-cmc

crates/store/re_types_core/src/loggable_batch.rs

@@ -44,6 +44,18 @@ pub trait ComponentBatch: LoggableBatch {
     /// Every component batch is uniquely identified by its [`ComponentDescriptor`].
     fn descriptor(&self) -> Cow<'_, ComponentDescriptor>;
 
+    // Wraps the current [`ComponentBatch`] with the given descriptor.
+    fn with_descriptor(
+        &self,
+        descriptor: ComponentDescriptor,
+    ) -> ComponentBatchCowWithDescriptor<'_>
+    where
+        Self: Sized,
+    {
+        ComponentBatchCowWithDescriptor::new(ComponentBatchCow::Ref(self as &dyn ComponentBatch))
+            .with_descriptor_override(descriptor)
+    }
+
     /// The fully-qualified name of this component batch, e.g. `rerun.components.Position2D`.
     ///
     /// This is a trivial but useful helper for `self.descriptor().component_name`.

docs/snippets/all/descriptors/descr_custom_archetype.py

@@ -5,50 +5,25 @@
 from typing import Any, Iterable
 
 import numpy.typing as npt
-import pyarrow as pa
 import rerun as rr  # pip install rerun-sdk
 
 
-class CustomPosition3DBatch(rr.ComponentBatchLike):
-    def __init__(self: Any, positions: npt.ArrayLike) -> None:
-        self.position = rr.components.Position3DBatch(positions)
-
-    def component_descriptor(self) -> rr.ComponentDescriptor:
-        return rr.ComponentDescriptor(
-            "user.CustomPosition3D",
-            archetype_name="user.CustomPoints3D",
-            archetype_field_name="custom_positions",
-        )
-
-    def as_arrow_array(self) -> pa.Array:
-        return self.position.as_arrow_array()
-
-
 class CustomPoints3D(rr.AsComponents):
     def __init__(self: Any, positions: npt.ArrayLike, colors: npt.ArrayLike) -> None:
-        self.positions = CustomPosition3DBatch(positions)
-        self.colors = rr.components.ColorBatch(colors)
+        self.positions = rr.components.Position3DBatch(positions).with_descriptor(
+            rr.ComponentDescriptor(
+                "user.CustomPosition3D",
+                archetype_name="user.CustomPoints3D",
+                archetype_field_name="custom_positions",
+            )
+        )
+        self.colors = rr.components.ColorBatch(colors).or_with_descriptor_overrides(
+            archetype_name="user.CustomPoints3D",
+            archetype_field_name="colors",
+        )
 
     def as_component_batches(self) -> Iterable[rr.ComponentBatchLike]:
-        return (
-            [rr.IndicatorComponentBatch("user.CustomPoints3D")]
-            + [
-                rr.DescribedComponentBatch(
-                    self.positions,
-                    self.positions.component_descriptor().or_with_overrides(
-                        archetype_name="user.CustomPoints3D", archetype_field_name="custom_positions"
-                    ),
-                )
-            ]
-            + [
-                rr.DescribedComponentBatch(
-                    self.colors,
-                    self.colors.component_descriptor().or_with_overrides(
-                        archetype_name="user.CustomPoints3D", archetype_field_name="colors"
-                    ),
-                )
-            ]
-        )
+        return [rr.IndicatorComponentBatch("user.CustomPoints3D"), self.positions, self.colors]
 
 
 rr.init("rerun_example_descriptors_custom_archetype")

docs/snippets/all/descriptors/descr_custom_archetype.rs

@@ -1,45 +1,7 @@
-use rerun::{
-    external::arrow2, ChunkStore, ChunkStoreConfig, Component, ComponentDescriptor, VersionPolicy,
-};
-
-#[derive(Debug, Clone, Copy)]
-struct CustomPosition3D(rerun::components::Position3D);
-
-impl rerun::SizeBytes for CustomPosition3D {
-    #[inline]
-    fn heap_size_bytes(&self) -> u64 {
-        0
-    }
-}
-
-impl rerun::Loggable for CustomPosition3D {
-    #[inline]
-    fn arrow2_datatype() -> arrow2::datatypes::DataType {
-        rerun::components::Position3D::arrow2_datatype()
-    }
-
-    #[inline]
-    fn to_arrow2_opt<'a>(
-        data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
-    ) -> rerun::SerializationResult<Box<dyn arrow2::array::Array>>
-    where
-        Self: 'a,
-    {
-        rerun::components::Position3D::to_arrow2_opt(
-            data.into_iter().map(|opt| opt.map(Into::into).map(|c| c.0)),
-        )
-    }
-}
-
-impl rerun::Component for CustomPosition3D {
-    #[inline]
-    fn descriptor() -> ComponentDescriptor {
-        ComponentDescriptor::new("user.CustomPosition3D")
-    }
-}
+use rerun::{ChunkStore, ChunkStoreConfig, ComponentBatch, ComponentDescriptor, VersionPolicy};
 
 struct CustomPoints3D {
-    positions: Vec<CustomPosition3D>,
+    positions: Vec<rerun::components::Position3D>,
     colors: Option<Vec<rerun::components::Color>>,
 }
 
@@ -49,13 +11,15 @@ impl CustomPoints3D {
     }
 
     fn overridden_position_descriptor() -> ComponentDescriptor {
-        CustomPosition3D::descriptor()
-            .or_with_archetype_name(|| "user.CustomPoints3D".into())
-            .or_with_archetype_field_name(|| "custom_positions".into())
+        ComponentDescriptor {
+            archetype_name: Some("user.CustomPoints3D".into()),
+            archetype_field_name: Some("custom_positions".into()),
+            component_name: "user.CustomPosition3D".into(),
+        }
     }
 
     fn overridden_color_descriptor() -> ComponentDescriptor {
-        rerun::components::Color::descriptor()
+        <rerun::components::Color as rerun::Component>::descriptor()
             .or_with_archetype_name(|| "user.CustomPoints3D".into())
             .or_with_archetype_field_name(|| "colors".into())
     }
@@ -66,15 +30,12 @@ impl rerun::AsComponents for CustomPoints3D {
         [
             Some(Self::indicator().to_batch()),
             Some(
-                rerun::ComponentBatchCowWithDescriptor::new(
-                    &self.positions as &dyn rerun::ComponentBatch,
-                )
-                .with_descriptor_override(Self::overridden_position_descriptor()),
+                self.positions
+                    .with_descriptor(Self::overridden_position_descriptor()),
             ),
-            self.colors.as_ref().map(|colors| {
-                rerun::ComponentBatchCowWithDescriptor::new(colors as &dyn rerun::ComponentBatch)
-                    .with_descriptor_override(Self::overridden_color_descriptor())
-            }),
+            self.colors
+                .as_ref()
+                .map(|colors| colors.with_descriptor(Self::overridden_color_descriptor())),
         ]
         .into_iter()
         .flatten()
@@ -83,12 +44,12 @@ impl rerun::AsComponents for CustomPoints3D {
 }
 
 fn example(rec: &rerun::RecordingStream) -> Result<(), Box<dyn std::error::Error>> {
-    let position = CustomPosition3D(rerun::components::Position3D::new(1.0, 2.0, 3.0));
-    let color = rerun::components::Color::new(0xFF00FFFF);
+    let positions = rerun::components::Position3D::new(1.0, 2.0, 3.0);
+    let colors = rerun::components::Color::new(0xFF00FFFF);
 
     let points = CustomPoints3D {
-        positions: vec![position],
-        colors: Some(vec![color]),
+        positions: vec![positions],
+        colors: Some(vec![colors]),
     };
 
     rec.log_static("data", &points as _)?;
@@ -152,7 +113,7 @@ fn check_tags(rec: &rerun::RecordingStream) {
             ComponentDescriptor {
                 archetype_name: Some("user.CustomPoints3D".into()),
                 archetype_field_name: Some("colors".into()),
-                component_name: rerun::components::Color::name(),
+                component_name: "rerun.components.Color".into(),
             },
             ComponentDescriptor {
                 archetype_name: Some("user.CustomPoints3D".into()),

docs/snippets/all/descriptors/descr_custom_component.py

@@ -2,31 +2,18 @@
 
 from __future__ import annotations
 
-from typing import Any
-
-import numpy.typing as npt
-import pyarrow as pa
 import rerun as rr  # pip install rerun-sdk
 
-
-class CustomPosition3DBatch(rr.ComponentBatchLike):
-    def __init__(self: Any, positions: npt.ArrayLike) -> None:
-        self.position = rr.components.Position3DBatch(positions)
-
-    def component_descriptor(self) -> rr.ComponentDescriptor:
-        return rr.ComponentDescriptor(
-            "user.CustomPosition3D",
-            archetype_name="user.CustomArchetype",
-            archetype_field_name="custom_positions",
-        )
-
-    def as_arrow_array(self) -> pa.Array:
-        return self.position.as_arrow_array()
-
-
 rr.init("rerun_example_descriptors_custom_component")
 rr.spawn()
 
-rr.log("data", [CustomPosition3DBatch([[1, 2, 3]])], static=True)
+positions = rr.components.Position3DBatch([1, 2, 3]).with_descriptor(
+    rr.ComponentDescriptor(
+        "user.CustomPosition3D",
+        archetype_name="user.CustomArchetype",
+        archetype_field_name="custom_positions",
+    )
+)
+rr.log("data", [positions], static=True)
 
 # The tags are indirectly checked by the Rust version (have a look over there for more info).

docs/snippets/all/descriptors/descr_custom_component.rs

@@ -1,51 +1,13 @@
-use rerun::{
-    external::arrow2, ChunkStore, ChunkStoreConfig, Component, ComponentDescriptor, Loggable,
-    VersionPolicy,
-};
-
-#[derive(Debug, Clone, Copy)]
-struct CustomPosition3D(rerun::components::Position3D);
-
-impl rerun::SizeBytes for CustomPosition3D {
-    #[inline]
-    fn heap_size_bytes(&self) -> u64 {
-        0
-    }
-}
-
-impl Loggable for CustomPosition3D {
-    #[inline]
-    fn arrow2_datatype() -> arrow2::datatypes::DataType {
-        rerun::components::Position3D::arrow2_datatype()
-    }
-
-    #[inline]
-    fn to_arrow2_opt<'a>(
-        data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
-    ) -> rerun::SerializationResult<Box<dyn arrow2::array::Array>>
-    where
-        Self: 'a,
-    {
-        rerun::components::Position3D::to_arrow2_opt(
-            data.into_iter().map(|opt| opt.map(Into::into).map(|c| c.0)),
-        )
-    }
-}
-
-impl Component for CustomPosition3D {
-    #[inline]
-    fn descriptor() -> ComponentDescriptor {
-        ComponentDescriptor {
-            archetype_name: Some("user.CustomArchetype".into()),
-            archetype_field_name: Some("custom_positions".into()),
-            component_name: "user.CustomPosition3D".into(),
-        }
-    }
-}
+use rerun::{ChunkStore, ChunkStoreConfig, ComponentBatch, ComponentDescriptor, VersionPolicy};
 
 fn example(rec: &rerun::RecordingStream) -> Result<(), Box<dyn std::error::Error>> {
-    let position = CustomPosition3D(rerun::components::Position3D::new(1.0, 2.0, 3.0));
-    rec.log_component_batches("data", true, [&position as &dyn rerun::ComponentBatch])?;
+    let positions = rerun::components::Position3D::new(1.0, 2.0, 3.0);
+    let positions = positions.with_descriptor(ComponentDescriptor {
+        archetype_name: Some("user.CustomArchetype".into()),
+        archetype_field_name: Some("custom_positions".into()),
+        component_name: "user.CustomPosition3D".into(),
+    });
+    rec.log_component_batches("data", true, [&positions as &dyn rerun::ComponentBatch])?;
 
     Ok(())
 }

docs/snippets/all/tutorials/custom_data.py

@@ -12,7 +12,7 @@
 import rerun as rr
 
 
-class ConfidenceBatch(rr.ComponentBatchLike):
+class ConfidenceBatch(rr.ComponentBatchMixin):
     """A batch of confidence data."""
 
     def __init__(self: Any, confidence: npt.ArrayLike) -> None:
@@ -30,24 +30,17 @@ def as_arrow_array(self) -> pa.Array:
 class CustomPoints3D(rr.AsComponents):
     """A custom archetype that extends Rerun's builtin `Points3D` archetype with a custom component."""
 
-    def __init__(self: Any, points3d: npt.ArrayLike, confidences: npt.ArrayLike) -> None:
-        self.points3d = points3d
-        self.confidences = confidences
+    def __init__(self: Any, positions: npt.ArrayLike, confidences: npt.ArrayLike) -> None:
+        self.points3d = rr.Points3D(positions)
+        self.confidences = ConfidenceBatch(confidences).or_with_descriptor_overrides(
+            archetype_name="user.CustomPoints3D", archetype_field_name="confidences"
+        )
 
     def as_component_batches(self) -> Iterable[rr.ComponentBatchLike]:
-        points3d = np.asarray(self.points3d)
-        confidences = ConfidenceBatch(self.confidences)
         return (
-            list(rr.Points3D(points3d).as_component_batches())  # The components from Points3D
+            list(self.points3d.as_component_batches())  # The components from Points3D
             + [rr.IndicatorComponentBatch("user.CustomPoints3D")]  # Our custom indicator
-            + [
-                rr.DescribedComponentBatch(
-                    confidences,
-                    confidences.component_descriptor().or_with_overrides(
-                        archetype_name="user.CustomPoints3D", archetype_field_name="confidences"
-                    ),
-                )
-            ]  # Custom confidence data
+            + [self.confidences]  # Custom confidence data
         )
 
 
@@ -59,14 +52,14 @@ def log_custom_data() -> None:
     rr.log(
         "left/my_confident_point_cloud",
         CustomPoints3D(
-            points3d=point_grid,
+            positions=point_grid,
             confidences=[42],
         ),
     )
 
     rr.log(
         "right/my_polarized_point_cloud",
-        CustomPoints3D(points3d=point_grid, confidences=np.arange(0, len(point_grid))),
+        CustomPoints3D(positions=point_grid, confidences=np.arange(0, len(point_grid))),
     )
 
 

rerun_py/rerun_sdk/rerun/__init__.py

@@ -33,6 +33,7 @@
     remote as remote,
 )
 from ._baseclasses import (
+    ComponentBatchMixin as ComponentBatchMixin,
     ComponentColumn as ComponentColumn,
     ComponentDescriptor as ComponentDescriptor,
     DescribedComponentBatch as DescribedComponentBatch,