Optimize gathering of point cloud colors

crates/re_arrow_store/benches/arrow2.rs

@@ -6,7 +6,7 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 use std::sync::Arc;
 
 use arrow2::array::{Array, PrimitiveArray, StructArray, UnionArray};
-use criterion::{criterion_group, Criterion};
+use criterion::Criterion;
 use itertools::Itertools;
 
 use re_log_types::{DataCell, SizeBytes as _};
@@ -19,7 +19,7 @@ use re_types::{
 
 // ---
 
-criterion_group!(benches, erased_clone, estimated_size_bytes);
+criterion::criterion_group!(benches, erased_clone, estimated_size_bytes);
 
 #[cfg(not(feature = "core_benchmarks_only"))]
 criterion::criterion_main!(benches);

crates/re_log_types/src/data_cell.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use arrow2::datatypes::DataType;
-use re_types::{Component, ComponentName, DeserializationError};
+use re_types::{Component, ComponentBatch, ComponentName, DeserializationError};
 
 use crate::SizeBytes;
 
@@ -164,6 +164,14 @@ pub struct DataCellInner {
 // TODO(#1696): Check that the array is indeed a leaf / component type when building a cell from an
 // arrow payload.
 impl DataCell {
+    /// Builds a new `DataCell` from a component batch.
+    #[inline]
+    pub fn from_component_batch(batch: &dyn ComponentBatch) -> re_types::SerializationResult<Self> {
+        batch
+            .to_arrow()
+            .map(|arrow| DataCell::from_arrow(batch.name(), arrow))
+    }
+
     /// Builds a new `DataCell` from a uniform iterable of native component values.
     ///
     /// Fails if the given iterable cannot be serialized to arrow, which should never happen when

crates/re_log_types/src/data_row.rs

@@ -1,6 +1,6 @@
 use ahash::HashSetExt;
 use nohash_hasher::IntSet;
-use re_types::ComponentName;
+use re_types::{AsComponents, ComponentName};
 use smallvec::SmallVec;
 
 use crate::{DataCell, DataCellError, DataTable, EntityPath, SizeBytes, TableId, TimePoint};
@@ -266,6 +266,48 @@ pub struct DataRow {
 }
 
 impl DataRow {
+    /// Builds a new `DataRow` from anything implementing [`AsComponents`].
+    pub fn from_archetype(
+        row_id: RowId,
+        timepoint: TimePoint,
+        entity_path: EntityPath,
+        as_components: &dyn AsComponents,
+    ) -> anyhow::Result<Self> {
+        re_tracing::profile_function!();
+
+        let batches = as_components.as_component_batches();
+        Self::from_component_batches(
+            row_id,
+            timepoint,
+            entity_path,
+            batches.iter().map(|batch| batch.as_ref()),
+        )
+    }
+
+    /// Builds a new `DataRow` from anything implementing [`AsComponents`].
+    pub fn from_component_batches<'a>(
+        row_id: RowId,
+        timepoint: TimePoint,
+        entity_path: EntityPath,
+        comp_batches: impl IntoIterator<Item = &'a dyn re_types::ComponentBatch>,
+    ) -> anyhow::Result<Self> {
+        re_tracing::profile_function!();
+
+        let data_cells = comp_batches
+            .into_iter()
+            .map(|batch| DataCell::from_component_batch(batch))
+            .collect::<Result<Vec<DataCell>, _>>()?;
+
+        // TODO(emilk): should `DataRow::from_cells` calculate `num_instances` instead?
+        let num_instances = data_cells.iter().map(|cell| cell.num_instances()).max();
+        let num_instances = num_instances.unwrap_or(0);
+
+        let mut row =
+            DataRow::from_cells(row_id, timepoint, entity_path, num_instances, data_cells)?;
+        row.compute_all_size_bytes();
+        Ok(row)
+    }
+
     /// Builds a new `DataRow` from an iterable of [`DataCell`]s.
     ///
     /// Fails if:

crates/re_space_view_spatial/Cargo.toml

@@ -44,3 +44,15 @@ parking_lot.workspace = true
 rayon.workspace = true
 serde = "1"
 smallvec = { workspace = true, features = ["serde"] }
+
+
+[dev-dependencies]
+criterion.workspace = true
+mimalloc.workspace = true
+
+[lib]
+bench = false
+
+[[bench]]
+name = "bench_points"
+harness = false

crates/re_space_view_spatial/benches/bench_points.rs

@@ -0,0 +1,134 @@
+//! High-level benchmark of the CPU-side of our `Points3D` rendering.
+
+use re_arrow_store::{DataStore, LatestAtQuery};
+use re_log_types::{DataRow, EntityPath, RowId, TimeInt, TimePoint, Timeline};
+use re_space_view_spatial::LoadedPoints;
+use re_types::{
+    archetypes::Points3D,
+    components::{Color, InstanceKey, Position3D},
+    Loggable as _,
+};
+use re_viewer_context::Annotations;
+
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+criterion::criterion_main!(benches);
+criterion::criterion_group!(benches, bench_points);
+
+// ---
+
+#[cfg(not(debug_assertions))]
+const NUM_POINTS: usize = 1_000_000;
+
+// `cargo test` also runs the benchmark setup code, so make sure they run quickly:
+#[cfg(debug_assertions)]
+const NUM_POINTS: usize = 10;
+
+// ---
+
+/// Mimics `examples/python/open_photogrammetry_format/main.py`
+fn bench_points(c: &mut criterion::Criterion) {
+    let timeline = Timeline::log_time();
+    let ent_path = EntityPath::from("points");
+
+    let store = {
+        let mut store = DataStore::new(InstanceKey::name(), Default::default());
+
+        let positions = vec![Position3D::new(0.1, 0.2, 0.3); NUM_POINTS];
+        let colors = vec![Color::from(0xffffffff); NUM_POINTS];
+        let points = Points3D::new(positions).with_colors(colors);
+        let mut timepoint = TimePoint::default();
+        timepoint.insert(timeline, TimeInt::from_seconds(0));
+        let data_row =
+            DataRow::from_archetype(RowId::random(), timepoint, ent_path.clone(), &points).unwrap();
+        store.insert_row(&data_row).unwrap();
+        store
+    };
+
+    let latest_at = LatestAtQuery::latest(timeline);
+    let annotations = Annotations::missing();
+
+    {
+        let mut group = c.benchmark_group("Points3D");
+        group.bench_function("query_archetype", |b| {
+            b.iter(|| {
+                let arch_view =
+                    re_query::query_archetype::<Points3D>(&store, &latest_at, &ent_path).unwrap();
+                assert_eq!(arch_view.num_instances(), NUM_POINTS);
+                arch_view
+            });
+        });
+    }
+
+    let arch_view = re_query::query_archetype::<Points3D>(&store, &latest_at, &ent_path).unwrap();
+    assert_eq!(arch_view.num_instances(), NUM_POINTS);
+
+    {
+        let mut group = c.benchmark_group("Points3D");
+        group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
+        group.bench_function("load_all", |b| {
+            b.iter(|| {
+                let points =
+                    LoadedPoints::load(&arch_view, &ent_path, latest_at.at, &annotations).unwrap();
+                assert_eq!(points.positions.len(), NUM_POINTS);
+                assert_eq!(points.colors.len(), NUM_POINTS);
+                assert_eq!(points.radii.len(), NUM_POINTS); // NOTE: we don't log radii, but we should get a list of defaults!
+                points
+            });
+        });
+    }
+
+    {
+        let mut group = c.benchmark_group("Points3D");
+        group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
+        group.bench_function("load_positions", |b| {
+            b.iter(|| {
+                let positions = LoadedPoints::load_positions(&arch_view).unwrap();
+                assert_eq!(positions.len(), NUM_POINTS);
+                positions
+            });
+        });
+    }
+
+    {
+        let points = LoadedPoints::load(&arch_view, &ent_path, latest_at.at, &annotations).unwrap();
+
+        let mut group = c.benchmark_group("Points3D");
+        group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
+        group.bench_function("load_colors", |b| {
+            b.iter(|| {
+                let colors =
+                    LoadedPoints::load_colors(&arch_view, &ent_path, &points.annotation_infos)
+                        .unwrap();
+                assert_eq!(colors.len(), NUM_POINTS);
+                colors
+            });
+        });
+    }
+
+    // NOTE: we don't log radii!
+    {
+        let mut group = c.benchmark_group("Points3D");
+        group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
+        group.bench_function("load_radii", |b| {
+            b.iter(|| {
+                let radii = LoadedPoints::load_radii(&arch_view, &ent_path).unwrap();
+                assert_eq!(radii.len(), NUM_POINTS);
+                radii
+            });
+        });
+    }
+
+    {
+        let mut group = c.benchmark_group("Points3D");
+        group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
+        group.bench_function("load_picking_ids", |b| {
+            b.iter(|| {
+                let picking_ids = LoadedPoints::load_picking_ids(&arch_view);
+                assert_eq!(picking_ids.len(), NUM_POINTS);
+                picking_ids
+            });
+        });
+    }
+}

crates/re_space_view_spatial/src/lib.rs

@@ -20,6 +20,9 @@ mod ui_3d;
 pub use space_view_2d::SpatialSpaceView2D;
 pub use space_view_3d::SpatialSpaceView3D;
 
+#[doc(hidden)] // Public for benchmarks
+pub use parts::LoadedPoints;
+
 // ---
 
 mod view_kind {

crates/re_space_view_spatial/src/parts/images.rs

@@ -259,10 +259,7 @@ impl ImagesPart {
                 .annotations
                 .resolved_class_description(None)
                 .annotation_info()
-                .color(
-                    color.map(|c| c.to_array()).as_ref(),
-                    DefaultColor::OpaqueWhite,
-                );
+                .color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite);
 
             if let Some(textured_rect) = to_textured_rect(
                 ctx,
@@ -379,10 +376,7 @@ impl ImagesPart {
                 .annotations
                 .resolved_class_description(None)
                 .annotation_info()
-                .color(
-                    color.map(|c| c.to_array()).as_ref(),
-                    DefaultColor::OpaqueWhite,
-                );
+                .color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite);
 
             if let Some(textured_rect) = to_textured_rect(
                 ctx,
@@ -467,10 +461,7 @@ impl ImagesPart {
                 .annotations
                 .resolved_class_description(None)
                 .annotation_info()
-                .color(
-                    color.map(|c| c.to_array()).as_ref(),
-                    DefaultColor::OpaqueWhite,
-                );
+                .color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite);
 
             if let Some(textured_rect) = to_textured_rect(
                 ctx,

crates/re_space_view_spatial/src/parts/mod.rs

@@ -22,8 +22,10 @@ use re_types::components::Text;
 pub use spatial_view_part::SpatialViewPartData;
 pub use transform3d_arrows::{add_axis_arrows, Transform3DArrowsPart};
 
+#[doc(hidden)] // Public for benchmarks
+pub use points3d::LoadedPoints;
+
 use ahash::HashMap;
-use std::sync::Arc;
 
 use re_data_store::{EntityPath, InstancePathHash};
 use re_types::components::{Color, InstanceKey};
@@ -118,7 +120,7 @@ pub fn process_colors<'a, A: Archetype>(
         arch_view.iter_optional_component::<Color>()?,
     )
     .map(move |(annotation_info, color)| {
-        annotation_info.color(color.map(move |c| c.to_array()).as_ref(), default_color)
+        annotation_info.color(color.map(|c| c.to_array()), default_color)
     }))
 }
 
@@ -134,9 +136,7 @@ pub fn process_labels<'a, A: Archetype>(
         annotation_infos.iter(),
         arch_view.iter_optional_component::<Text>()?,
     )
-    .map(move |(annotation_info, text)| {
-        annotation_info.label(text.as_ref().map(move |t| t.as_str()))
-    }))
+    .map(move |(annotation_info, text)| annotation_info.label(text.as_ref().map(|t| t.as_str()))))
 }
 
 /// Process [`re_types::components::Radius`] components to [`re_renderer::Size`] using auto size
@@ -171,22 +171,22 @@ pub fn process_radii<'a, A: Archetype>(
 fn process_annotations<Primary, A: Archetype>(
     query: &ViewQuery<'_>,
     arch_view: &re_query::ArchetypeView<A>,
-    annotations: &Arc<Annotations>,
+    annotations: &Annotations,
 ) -> Result<ResolvedAnnotationInfos, re_query::QueryError>
 where
     Primary: re_types::Component + Clone,
 {
-    process_annotations_and_keypoints(query, arch_view, annotations, |_: &Primary| {
+    process_annotations_and_keypoints(query.latest_at, arch_view, annotations, |_: &Primary| {
         glam::Vec3::ZERO
     })
     .map(|(a, _)| a)
 }
 
 /// Resolves all annotations and keypoints for the given entity view.
 fn process_annotations_and_keypoints<Primary, A: Archetype>(
-    query: &ViewQuery<'_>,
+    latest_at: re_log_types::TimeInt,
     arch_view: &re_query::ArchetypeView<A>,
-    annotations: &Arc<Annotations>,
+    annotations: &Annotations,
     mut primary_into_position: impl FnMut(&Primary) -> glam::Vec3,
 ) -> Result<(ResolvedAnnotationInfos, Keypoints), re_query::QueryError>
 where
@@ -220,7 +220,7 @@ where
 
         if let (Some(keypoint_id), Some(class_id), primary) = (keypoint_id, class_id, primary) {
             keypoints
-                .entry((class_id, query.latest_at.as_i64()))
+                .entry((class_id, latest_at.as_i64()))
                 .or_default()
                 .insert(keypoint_id.0, primary_into_position(&primary));
             class_description.annotation_info_with_keypoint(keypoint_id.0)