Centralize freestanding store helpers (#2153)

* centralize freestanding store helpers

* port everything to new helpers

* note to self

* missing feature flag

* docs

crates/re_arrow_store/src/lib.rs

@@ -20,6 +20,7 @@ mod store_arrow;
 mod store_dump;
 mod store_format;
 mod store_gc;
+mod store_helpers;
 mod store_read;
 mod store_sanity;
 mod store_stats;

crates/re_arrow_store/src/store_helpers.rs

@@ -0,0 +1,158 @@
+use re_log_types::{
+    ComponentName, DataCell, DataRow, DeserializableComponent, EntityPath, RowId,
+    SerializableComponent, TimeInt, TimePoint, Timeline,
+};
+
+use crate::{DataStore, LatestAtQuery};
+
+// --- Read ---
+
+impl DataStore {
+    /// Get the latest value for a given [`re_log_types::Component`].
+    ///
+    /// This assumes that the row we get from the store only contains a single instance for this
+    /// component; it will log a warning otherwise.
+    ///
+    /// This should only be used for "mono-components" such as `Transform` and `Tensor`.
+    ///
+    /// This is a best-effort helper, it will merely log errors on failure.
+    pub fn query_latest_component<C: DeserializableComponent>(
+        &self,
+        entity_path: &EntityPath,
+        query: &LatestAtQuery,
+    ) -> Option<C>
+    where
+        for<'b> &'b C::ArrayType: IntoIterator,
+    {
+        crate::profile_function!();
+
+        let (_, cells) = self.latest_at(query, entity_path, C::name(), &[C::name()])?;
+        let cell = cells.get(0)?.as_ref()?;
+
+        let mut iter = cell
+            .try_to_native::<C>()
+            .map_err(|err| {
+                re_log::error_once!(
+                    "Couldn't deserialize component at {entity_path}.{}: {err}",
+                    C::name()
+                );
+            })
+            .ok()?;
+
+        let component = iter.next();
+
+        if iter.next().is_some() {
+            re_log::warn_once!("Unexpected batch for {} at: {}", C::name(), entity_path);
+        }
+
+        component
+    }
+
+    /// Get the latest value for a given [`re_log_types::Component`], assuming it is timeless.
+    ///
+    /// This assumes that the row we get from the store only contains a single instance for this
+    /// component; it will log a warning otherwise.
+    ///
+    /// This should only be used for "mono-components" such as `Transform` and `Tensor`.
+    ///
+    /// This is a best-effort helper, it will merely log errors on failure.
+    pub fn query_timeless_component<C: DeserializableComponent>(
+        &self,
+        entity_path: &EntityPath,
+    ) -> Option<C>
+    where
+        for<'b> &'b C::ArrayType: IntoIterator,
+    {
+        crate::profile_function!();
+
+        let query = LatestAtQuery::new(Timeline::default(), TimeInt::MAX);
+        self.query_latest_component(entity_path, &query)
+    }
+}
+
+// --- Write ---
+
+impl DataStore {
+    /// Stores a single value for a given [`re_log_types::Component`].
+    ///
+    /// This is a best-effort helper, it will merely log errors on failure.
+    pub fn insert_component<C: SerializableComponent>(
+        &mut self,
+        entity_path: &EntityPath,
+        timepoint: &TimePoint,
+        component: C,
+    ) {
+        crate::profile_function!();
+
+        let mut row = match DataRow::try_from_cells1(
+            RowId::random(),
+            entity_path.clone(),
+            timepoint.clone(),
+            1,
+            [component].as_slice(),
+        ) {
+            Ok(row) => row,
+            Err(err) => {
+                re_log::error_once!(
+                    "Couldn't serialize component at {entity_path}.{}: {err}",
+                    C::name()
+                );
+                return;
+            }
+        };
+        row.compute_all_size_bytes();
+
+        if let Err(err) = self.insert_row(&row) {
+            re_log::error_once!(
+                "Couldn't insert component at {entity_path}.{}: {err}",
+                C::name()
+            );
+        }
+    }
+
+    /// Stores a single empty value for a given [`re_log_types::ComponentName`].
+    ///
+    /// This is a best-effort helper, it will merely log errors on failure.
+    pub fn insert_empty_component(
+        &mut self,
+        entity_path: &EntityPath,
+        timepoint: &TimePoint,
+        component: ComponentName,
+    ) {
+        crate::profile_function!();
+
+        if let Some(datatype) = self.lookup_datatype(&component) {
+            let cell = DataCell::from_arrow_empty(component, datatype.clone());
+
+            let mut row = match DataRow::try_from_cells1(
+                RowId::random(),
+                entity_path.clone(),
+                timepoint.clone(),
+                cell.num_instances(),
+                cell,
+            ) {
+                Ok(row) => row,
+                Err(err) => {
+                    re_log::error_once!(
+                        "Couldn't serialize component at {entity_path}.{}: {err}",
+                        component
+                    );
+                    return;
+                }
+            };
+            row.compute_all_size_bytes();
+
+            if let Err(err) = self.insert_row(&row) {
+                re_log::error_once!(
+                    "Couldn't insert component at {entity_path}.{}: {err}",
+                    component
+                );
+            }
+        } else {
+            re_log::error_once!(
+                "Couldn't find appropriate datatype at {entity_path}.{}",
+                component
+            );
+        }
+    }
+}

crates/re_data_store/src/log_db.rs

@@ -88,6 +88,8 @@ impl EntityDb {
                 let cell =
                     DataCell::from_arrow_empty(cell.component_name(), cell.datatype().clone());
 
+                // NOTE(cmc): The fact that this inserts data to multiple entity paths using a
+                // single `RowId` is... interesting. Keep it in mind.
                 let row = DataRow::from_cells1(
                     row_id,
                     row.entity_path.clone(),
@@ -117,6 +119,9 @@ impl EntityDb {
                 // TODO(jleibs): Faster empty-array creation
                 let cell =
                     DataCell::from_arrow_empty(component_path.component_name, data_type.clone());
+
+                // NOTE(cmc): The fact that this inserts data to multiple entity paths using a
+                // single `RowId` is... interesting. Keep it in mind.
                 let row = DataRow::from_cells1(
                     row_id,
                     component_path.entity_path.clone(),

crates/re_data_ui/src/annotation_context.rs

@@ -1,7 +1,7 @@
 use egui::{color_picker, Vec2};
 use itertools::Itertools;
 
-use re_log_types::{context::AnnotationInfo, AnnotationContext};
+use re_log_types::{component_types::ClassId, context::AnnotationInfo, AnnotationContext};
 use re_viewer_context::{auto_color, UiVerbosity, ViewerContext};
 
 use super::DataUi;
@@ -77,8 +77,11 @@ fn annotation_info(
     query: &re_arrow_store::LatestAtQuery,
     keypoint_id: &re_log_types::component_types::KeypointId,
 ) -> Option<re_log_types::context::AnnotationInfo> {
-    let class_id =
-        re_data_store::query_latest_single(&ctx.log_db.entity_db.data_store, entity_path, query)?;
+    let class_id = ctx
+        .log_db
+        .entity_db
+        .data_store
+        .query_latest_component::<ClassId>(entity_path, query)?;
     let annotations = crate::annotations(ctx, query, entity_path);
     let class = annotations
         .class_description(Some(class_id))

crates/re_query/src/query.rs

@@ -79,7 +79,7 @@ pub fn get_component_with_instances(
 /// length.
 ///
 /// If you expect only one instance (e.g. for mono-components like `Transform` `Tensor`]
-/// and have no additional components you can use [`re_data_store::query_latest_single`] instead.
+/// and have no additional components you can use [`DataStore::query_latest_component`] instead.
 ///
 /// ```
 /// # use re_arrow_store::LatestAtQuery;

crates/re_viewer/src/misc/queries.rs

@@ -1,5 +1,5 @@
 use re_arrow_store::LatestAtQuery;
-use re_data_store::{query_latest_single, EntityPath};
+use re_data_store::EntityPath;
 use re_log_types::Transform;
 use re_viewer_context::ViewerContext;
 
@@ -11,13 +11,12 @@ pub fn closest_pinhole_transform(
 ) -> Option<EntityPath> {
     crate::profile_function!();
 
-    let data_store = &ctx.log_db.entity_db.data_store;
+    let store = &ctx.log_db.entity_db.data_store;
 
     let mut pinhole_ent_path = None;
     let mut cur_path = Some(entity_path.clone());
     while let Some(path) = cur_path {
-        if let Some(Transform::Pinhole(_)) =
-            query_latest_single::<Transform>(data_store, &path, query)
+        if let Some(Transform::Pinhole(_)) = store.query_latest_component::<Transform>(&path, query)
         {
             pinhole_ent_path = Some(path);
             break;

crates/re_viewer/src/misc/space_info.rs

@@ -3,7 +3,7 @@ use std::collections::BTreeMap;
 use nohash_hasher::IntSet;
 
 use re_arrow_store::{LatestAtQuery, TimeInt, Timeline};
-use re_data_store::{log_db::EntityDb, query_latest_single, EntityPath, EntityTree};
+use re_data_store::{log_db::EntityDb, EntityPath, EntityTree};
 use re_log_types::Transform;
 
 use super::UnreachableTransform;
@@ -111,8 +111,9 @@ impl SpaceInfoCollection {
             tree: &EntityTree,
             query: &LatestAtQuery,
         ) {
-            if let Some(transform) =
-                query_latest_single::<Transform>(&entity_db.data_store, &tree.path, query)
+            if let Some(transform) = entity_db
+                .data_store
+                .query_latest_component::<Transform>(&tree.path, query)
             {
                 // A set transform (likely non-identity) - create a new space.
                 parent_space
@@ -157,7 +158,9 @@ impl SpaceInfoCollection {
         let mut spaces_info = Self::default();
 
         // Start at the root. The root is always part of the collection!
-        if query_latest_single::<Transform>(&entity_db.data_store, &EntityPath::root(), &query)
+        if entity_db
+            .data_store
+            .query_latest_component::<Transform>(&EntityPath::root(), &query)
             .is_some()
         {
             re_log::warn_once!("The root entity has a 'transform' component! This will have no effect. Did you mean to apply the transform elsewhere?");

crates/re_viewer/src/misc/transform_cache.rs

@@ -1,8 +1,6 @@
 use nohash_hasher::IntMap;
 use re_arrow_store::LatestAtQuery;
-use re_data_store::{
-    log_db::EntityDb, query_latest_single, EntityPath, EntityPropertyMap, EntityTree,
-};
+use re_data_store::{log_db::EntityDb, EntityPath, EntityPropertyMap, EntityTree};
 use re_viewer_context::TimeControl;
 
 /// Provides transforms from an entity to a chosen reference space for all elements in the scene
@@ -224,12 +222,12 @@ impl TransformCache {
 
 fn transform_at(
     entity_path: &EntityPath,
-    data_store: &re_arrow_store::DataStore,
+    store: &re_arrow_store::DataStore,
     query: &LatestAtQuery,
     pinhole_image_plane_distance: impl Fn(&EntityPath) -> f32,
     encountered_pinhole: &mut bool,
 ) -> Result<Option<glam::Affine3A>, UnreachableTransform> {
-    if let Some(transform) = query_latest_single(data_store, entity_path, query) {
+    if let Some(transform) = store.query_latest_component(entity_path, query) {
         match transform {
             re_log_types::Transform::Rigid3(rigid) => Ok(Some(rigid.parent_from_child().into())),
             // If we're connected via 'unknown' it's not reachable

crates/re_viewer/src/ui/selection_panel.rs

@@ -1,8 +1,6 @@
 use egui::NumExt as _;
 
-use re_data_store::{
-    query_latest_single, ColorMapper, Colormap, EditableAutoValue, EntityPath, EntityProperties,
-};
+use re_data_store::{ColorMapper, Colormap, EditableAutoValue, EntityPath, EntityProperties};
 use re_data_ui::{item_ui, DataUi};
 use re_log_types::{
     component_types::{Tensor, TensorDataMeaning},
@@ -439,8 +437,9 @@ fn pinhole_props_ui(
     entity_props: &mut EntityProperties,
 ) {
     let query = ctx.current_query();
+    let store = &ctx.log_db.entity_db.data_store;
     if let Some(re_log_types::Transform::Pinhole(_)) =
-        query_latest_single::<Transform>(&ctx.log_db.entity_db.data_store, entity_path, &query)
+        store.query_latest_component::<Transform>(entity_path, &query)
     {
         ui.label("Image plane distance");
         let mut distance = *entity_props.pinhole_image_plane_distance.get();
@@ -469,8 +468,8 @@ fn depth_props_ui(
     crate::profile_function!();
 
     let query = ctx.current_query();
-    let tensor =
-        query_latest_single::<Tensor>(&ctx.log_db.entity_db.data_store, entity_path, &query)?;
+    let store = &ctx.log_db.entity_db.data_store;
+    let tensor = store.query_latest_component::<Tensor>(entity_path, &query)?;
     if tensor.meaning != TensorDataMeaning::Depth {
         return Some(());
     }