Speed up Point3D (#1064)

* Don't process annotations if the components aren't presetn
* Pre-allocate storage for aggregating points, colors, and user-data

crates/re_query/src/entity_view.rs

@@ -312,6 +312,13 @@ where
         self.primary.iter_values()
     }
 
+    /// Check if the entity has a component and its not empty
+    pub fn has_component<C: Component>(&self) -> bool {
+        self.components
+            .get(&C::name())
+            .map_or(false, |c| !c.is_empty())
+    }
+
     /// Iterate over the values of a `Component`.
     ///
     /// Always produces an iterator of length `self.primary.len()`

crates/re_renderer/examples/2d.rs

@@ -153,6 +153,7 @@ impl framework::Example for Render2D {
         point_cloud_builder
             .batch("points")
             .add_points_2d(
+                4,
                 [
                     glam::vec2(500.0, 120.0),
                     glam::vec2(520.0, 120.0),
@@ -170,7 +171,7 @@ impl framework::Example for Render2D {
                 ]
                 .into_iter(),
             )
-            .color(Color32::from_rgb(55, 180, 1));
+            .colors(std::iter::repeat(Color32::from_rgb(55, 180, 1)).take(4));
 
         // Pile stuff to test for overlap handling
         {

crates/re_renderer/src/point_cloud_builder.rs

@@ -35,7 +35,7 @@ impl<PerPointUserData> Default for PointCloudBuilder<PerPointUserData> {
 
 impl<PerPointUserData> PointCloudBuilder<PerPointUserData>
 where
-    PerPointUserData: Default + Clone,
+    PerPointUserData: Default + Copy,
 {
     /// Start of a new batch.
     #[inline]
@@ -111,14 +111,20 @@ where
 
 pub struct PointCloudBatchBuilder<'a, PerPointUserData>(
     &'a mut PointCloudBuilder<PerPointUserData>,
-);
+)
+where
+    PerPointUserData: Default + Copy;
 
-impl<'a, PerPointUserData> Drop for PointCloudBatchBuilder<'a, PerPointUserData> {
+impl<'a, PerPointUserData> Drop for PointCloudBatchBuilder<'a, PerPointUserData>
+where
+    PerPointUserData: Default + Copy,
+{
     fn drop(&mut self) {
         // Remove batch again if it wasn't actually used.
         if self.0.batches.last().unwrap().point_count == 0 {
             self.0.batches.pop();
         }
+        self.extend_defaults();
     }
 }
 
@@ -141,18 +147,51 @@ where
         self
     }
 
+    /// Each time we `add_points`, or upon builder drop, make sure that we
+    /// fill in any additional colors and user-data to have matched vectors.
+    fn extend_defaults(&mut self) {
+        if self.0.colors.len() < self.0.vertices.len() {
+            self.0.colors.extend(
+                std::iter::repeat(Color32::WHITE).take(self.0.vertices.len() - self.0.colors.len()),
+            );
+        }
+
+        if self.0.user_data.len() < self.0.user_data.len() {
+            self.0.user_data.extend(
+                std::iter::repeat(PerPointUserData::default())
+                    .take(self.0.vertices.len() - self.0.user_data.len()),
+            );
+        }
+    }
+
     #[inline]
+    /// Add several 3D points
+    ///
+    /// Returns a `PointBuilder` which can be used to set the colors, radii, and user-data for the points.
+    ///
+    /// Params:
+    ///  - `size_hint`: The `PointBuilder` will pre-allocate buffers to accomodate up to this number of points.
+    ///                 The resulting point batch, will still be determined by the length of the iterator.
+    ///  - `positions`: An iterable of the positions of the collection of points
     pub fn add_points(
         &mut self,
+        size_hint: usize,
         positions: impl Iterator<Item = glam::Vec3>,
     ) -> PointsBuilder<'_, PerPointUserData> {
+        // TODO(jleibs): Figure out if we can plumb-through proper support for `Iterator::size_hints()`
+        // or potentially make `FixedSizedIterator` work correctly. This should be possible size the
+        // underlying arrow structures are of known-size, but carries some complexity with the amount of
+        // chaining, joining, filtering, etc. that happens along the way.
         crate::profile_function!();
 
+        self.extend_defaults();
+
         debug_assert_eq!(self.0.vertices.len(), self.0.colors.len());
         debug_assert_eq!(self.0.vertices.len(), self.0.user_data.len());
 
         let old_size = self.0.vertices.len();
 
+        self.0.vertices.reserve(size_hint);
         self.0.vertices.extend(positions.map(|p| PointCloudVertex {
             position: p,
             radius: Size::AUTO,
@@ -161,19 +200,18 @@ where
         let num_points = self.0.vertices.len() - old_size;
         self.batch_mut().point_count += num_points as u32;
 
-        self.0
-            .colors
-            .extend(std::iter::repeat(Color32::WHITE).take(num_points));
-        self.0
-            .user_data
-            .extend(std::iter::repeat(PerPointUserData::default()).take(num_points));
+        self.0.colors.reserve(num_points);
+        self.0.user_data.reserve(num_points);
 
         let new_range = old_size..self.0.vertices.len();
 
+        let max_points = self.0.vertices.len();
+
         PointsBuilder {
-            vertices: &mut self.0.vertices[new_range.clone()],
-            colors: &mut self.0.colors[new_range.clone()],
-            user_data: &mut self.0.user_data[new_range],
+            vertices: &mut self.0.vertices[new_range],
+            max_points,
+            colors: &mut self.0.colors,
+            user_data: &mut self.0.user_data,
         }
     }
 
@@ -198,12 +236,18 @@ where
     }
 
     /// Adds several 2D points. Uses an autogenerated depth value, the same for all points passed.
+    ///
+    /// Params:
+    ///  - `size_hint`: The `PointBuilder` will pre-allocate buffers to accomodate up to this number of points.
+    ///                 The resulting point batch, will be the size of the length of the `positions` iterator.
+    ///  - `positions`: An iterable of the positions of the collection of points
     #[inline]
     pub fn add_points_2d(
         &mut self,
+        size_hint: usize,
         positions: impl Iterator<Item = glam::Vec2>,
     ) -> PointsBuilder<'_, PerPointUserData> {
-        self.add_points(positions.map(|p| p.extend(0.0)))
+        self.add_points(size_hint, positions.map(|p| p.extend(0.0)))
     }
 
     /// Adds a single 2D point. Uses an autogenerated depth value.
@@ -248,67 +292,44 @@ where
 }
 
 pub struct PointsBuilder<'a, PerPointUserData> {
+    // Vertices is a slice, which radii will update
     vertices: &'a mut [PointCloudVertex],
-    colors: &'a mut [Color32],
-    user_data: &'a mut [PerPointUserData],
+
+    // Colors and user-data are the Vec we append
+    // the data to if provided.
+    max_points: usize,
+    colors: &'a mut Vec<Color32>,
+    user_data: &'a mut Vec<PerPointUserData>,
 }
 
 impl<'a, PerPointUserData> PointsBuilder<'a, PerPointUserData>
 where
     PerPointUserData: Clone,
 {
-    /// Splats a radius to all points in this builder.
-    #[inline]
-    pub fn radius(self, radius: Size) -> Self {
-        for point in self.vertices.iter_mut() {
-            point.radius = radius;
-        }
-        self
-    }
-
     /// Assigns radii to all points.
     ///
     /// If the iterator doesn't cover all points, some will not be assigned.
     /// If the iterator provides more values than there are points, the extra values will be ignored.
     #[inline]
     pub fn radii(self, radii: impl Iterator<Item = Size>) -> Self {
+        // TODO(andreas): This seems like an argument for moving radius
+        // to a separate storage
         crate::profile_function!();
         for (point, radius) in self.vertices.iter_mut().zip(radii) {
             point.radius = radius;
         }
         self
     }
 
-    /// Splats a color to all points in this builder.
-    #[inline]
-    pub fn color(self, color: Color32) -> Self {
-        for c in self.colors.iter_mut() {
-            *c = color;
-        }
-        self
-    }
-
     /// Assigns colors to all points.
     ///
     /// If the iterator doesn't cover all points, some will not be assigned.
     /// If the iterator provides more values than there are points, the extra values will be ignored.
     #[inline]
     pub fn colors(self, colors: impl Iterator<Item = Color32>) -> Self {
         crate::profile_function!();
-        for (c, color) in self.colors.iter_mut().zip(colors) {
-            *c = color;
-        }
-        self
-    }
-
-    /// Splats a single user data to all points in this builder.
-    ///
-    /// User data is currently not available on the GPU.
-    #[inline]
-    pub fn user_data_splat(self, data: PerPointUserData) -> Self {
-        for user_data in self.user_data.iter_mut() {
-            *user_data = data.clone();
-        }
+        self.colors
+            .extend(colors.take(self.max_points - self.colors.len()));
         self
     }
 
@@ -318,9 +339,8 @@ where
     #[inline]
     pub fn user_data(self, data: impl Iterator<Item = PerPointUserData>) -> Self {
         crate::profile_function!();
-        for (d, data) in self.user_data.iter_mut().zip(data) {
-            *d = data;
-        }
+        self.user_data
+            .extend(data.take(self.max_points - self.user_data.len()));
         self
     }
 }

crates/re_viewer/src/ui/annotations.rs

@@ -66,6 +66,7 @@ pub enum DefaultColor<'a> {
     EntityPath(&'a EntityPath),
 }
 
+#[derive(Clone)]
 pub struct ResolvedAnnotationInfo(pub Option<AnnotationInfo>);
 
 impl ResolvedAnnotationInfo {

crates/re_viewer/src/ui/view_spatial/scene/scene_part/points3d.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use ahash::{HashMap, HashMapExt};
-use glam::{Mat4, Vec3};
+use glam::Mat4;
 
 use re_data_store::{EntityPath, EntityProperties};
 use re_log_types::{
@@ -39,27 +39,35 @@ impl Points3DPart {
         query: &SceneQuery<'_>,
         entity_view: &EntityView<Point3D>,
         annotations: &Arc<Annotations>,
-        point_positions: &[Vec3],
     ) -> Result<(Vec<ResolvedAnnotationInfo>, Keypoints), QueryError> {
         crate::profile_function!();
 
         let mut keypoints: Keypoints = HashMap::new();
 
+        // No need to process annotations if we don't have keypoints or class-ids
+        if !entity_view.has_component::<KeypointId>() && !entity_view.has_component::<ClassId>() {
+            let resolved_annotation = annotations.class_description(None).annotation_info();
+            return Ok((
+                vec![resolved_annotation; entity_view.num_instances()],
+                keypoints,
+            ));
+        }
+
         let annotation_info = itertools::izip!(
-            point_positions.iter(),
+            entity_view.iter_primary()?,
             entity_view.iter_component::<KeypointId>()?,
             entity_view.iter_component::<ClassId>()?,
         )
         .map(|(position, keypoint_id, class_id)| {
             let class_description = annotations.class_description(class_id);
 
-            if let Some(keypoint_id) = keypoint_id {
-                if let Some(class_id) = class_id {
-                    keypoints
-                        .entry((class_id, query.latest_at.as_i64()))
-                        .or_insert_with(Default::default)
-                        .insert(keypoint_id, *position);
-                }
+            if let (Some(keypoint_id), Some(class_id), Some(position)) =
+                (keypoint_id, class_id, position)
+            {
+                keypoints
+                    .entry((class_id, query.latest_at.as_i64()))
+                    .or_insert_with(Default::default)
+                    .insert(keypoint_id, position.into());
                 class_description.annotation_info_with_keypoint(keypoint_id)
             } else {
                 class_description.annotation_info()
@@ -153,15 +161,10 @@ impl Points3DPart {
             entity_view
                 .iter_primary()?
                 .filter_map(|pt| pt.map(glam::Vec3::from))
-                .collect::<Vec<_>>()
         };
 
-        let (annotation_infos, keypoints) = Self::process_annotations(
-            query,
-            entity_view,
-            &annotations,
-            point_positions.as_slice(),
-        )?;
+        let (annotation_infos, keypoints) =
+            Self::process_annotations(query, entity_view, &annotations)?;
         let instance_path_hashes = {
             crate::profile_scope!("instance_hashes");
             entity_view
@@ -200,7 +203,7 @@ impl Points3DPart {
             .points
             .batch("3d points")
             .world_from_obj(world_from_obj)
-            .add_points(point_positions.into_iter())
+            .add_points(entity_view.num_instances(), point_positions)
             .colors(colors)
             .radii(radii)
             .user_data(instance_path_hashes.into_iter());