New data APIs 5: port data UIs to new APIs (#5633)

_Trivial commit by commit_

All data UIs are now using cached APIs as much as possible.

As a nice side-effect, `EntityDb` is now plumbed through everywhere,
giving access to everything you might possibly need in all the places
you might need 'em.

We're slowly but surely seeing the first sign of instance keys going
away, but it's still not the focus of this PR.

---

Part of a PR series to completely revamp the data APIs in preparation
for the removal of instance keys and the introduction of promises:
- #5573
- #5574
- #5581
- #5605
- #5606
- #5633
- #5673
- #5679
- #5687
- #5755
- TODO
- TODO

Builds on top of the static data PR series:
- #5534

crates/re_data_ui/src/annotation_context.rs

@@ -17,7 +17,7 @@ impl crate::EntityDataUi for re_types::components::ClassId {
         verbosity: re_viewer_context::UiVerbosity,
         entity_path: &re_log_types::EntityPath,
         query: &re_data_store::LatestAtQuery,
-        _store: &re_data_store::DataStore,
+        _db: &re_entity_db::EntityDb,
     ) {
         let annotations = crate::annotations(ctx, query, entity_path);
         let class = annotations
@@ -63,7 +63,7 @@ impl crate::EntityDataUi for re_types::components::KeypointId {
         _verbosity: re_viewer_context::UiVerbosity,
         entity_path: &re_log_types::EntityPath,
         query: &re_data_store::LatestAtQuery,
-        _store: &re_data_store::DataStore,
+        _db: &re_entity_db::EntityDb,
     ) {
         if let Some(info) = annotation_info(ctx, entity_path, query, self.0) {
             ui.horizontal(|ui| {
@@ -102,7 +102,7 @@ impl DataUi for AnnotationContext {
         ui: &mut egui::Ui,
         verbosity: UiVerbosity,
         _query: &re_data_store::LatestAtQuery,
-        _store: &re_data_store::DataStore,
+        _db: &re_entity_db::EntityDb,
     ) {
         match verbosity {
             UiVerbosity::Small | UiVerbosity::Reduced => {

crates/re_data_ui/src/app_id.rs

@@ -13,7 +13,7 @@ impl crate::DataUi for ApplicationId {
         ui: &mut egui::Ui,
         verbosity: UiVerbosity,
         _query: &re_data_store::LatestAtQuery,
-        _store: &re_data_store::DataStore,
+        _db: &EntityDb,
     ) {
         egui::Grid::new("application_id")
             .num_columns(2)

crates/re_data_ui/src/blueprint_data.rs

@@ -10,12 +10,12 @@ impl<T: BlueprintIdRegistry> DataUi for BlueprintId<T> {
         ui: &mut egui::Ui,
         _verbosity: UiVerbosity,
         query: &re_data_store::LatestAtQuery,
-        store: &re_data_store::DataStore,
+        db: &re_entity_db::EntityDb,
     ) {
         entity_path_button_to(
             ctx,
             query,
-            store,
+            db,
             ui,
             None,
             &self.as_entity_path(),

crates/re_data_ui/src/component.rs

@@ -1,46 +1,44 @@
+use std::sync::Arc;
+
 use egui::NumExt;
 
-use re_entity_db::{EntityPath, InstancePath};
-use re_query::ComponentWithInstances;
+use re_entity_db::{
+    external::re_query_cache2::CachedLatestAtComponentResults, EntityPath, InstancePath,
+};
 use re_types::ComponentName;
 use re_ui::SyntaxHighlighting as _;
 use re_viewer_context::{UiVerbosity, ViewerContext};
 
 use super::{table_for_verbosity, DataUi};
 use crate::item_ui;
 
-// We do NOT implement `DataUi` for just `ComponentWithInstances`
-// because we also want the context of what entity it is part of!
-
 /// All the values of a specific [`re_log_types::ComponentPath`].
-pub struct EntityComponentWithInstances {
+pub struct EntityLatestAtResults {
     pub entity_path: EntityPath,
-    pub component_data: ComponentWithInstances,
-}
-
-impl EntityComponentWithInstances {
-    pub fn component_name(&self) -> ComponentName {
-        self.component_data.name()
-    }
-
-    pub fn num_instances(&self) -> usize {
-        self.component_data.len()
-    }
+    pub component_name: ComponentName,
+    pub results: Arc<CachedLatestAtComponentResults>,
 }
 
-impl DataUi for EntityComponentWithInstances {
+impl DataUi for EntityLatestAtResults {
     fn data_ui(
         &self,
         ctx: &ViewerContext<'_>,
         ui: &mut egui::Ui,
         verbosity: UiVerbosity,
         query: &re_data_store::LatestAtQuery,
-        store: &re_data_store::DataStore,
+        db: &re_entity_db::EntityDb,
     ) {
-        re_tracing::profile_function!(self.component_name().full_name());
+        re_tracing::profile_function!(self.component_name);
 
-        let instance_keys = self.component_data.instance_keys();
-        let num_instances = self.num_instances();
+        // TODO(#5607): what should happen if the promise is still pending?
+        let Some(num_instances) = self
+            .results
+            .raw(db.resolver(), self.component_name)
+            .map(|data| data.len())
+        else {
+            ui.weak("<pending>");
+            return;
+        };
 
         let one_line = match verbosity {
             UiVerbosity::Small => true,
@@ -71,7 +69,7 @@ impl DataUi for EntityComponentWithInstances {
         // ├───┼───┼───┼───┤ │
         // │   │ x │…+2│…+3│ │
         // └───┴───┴───┴───┘ ┘
-        let displayed_row = if num_instances <= max_row {
+        let num_displayed_rows = if num_instances <= max_row {
             num_instances
         } else {
             // this accounts for the "…x more" using a row and handles `num_instances == 0`
@@ -81,20 +79,16 @@ impl DataUi for EntityComponentWithInstances {
         if num_instances == 0 {
             ui.weak("(empty)");
         } else if num_instances == 1 {
-            if let Some(instance_key) = instance_keys.first() {
-                ctx.component_ui_registry.ui(
-                    ctx,
-                    ui,
-                    verbosity,
-                    query,
-                    store,
-                    &self.entity_path,
-                    &self.component_data,
-                    instance_key,
-                );
-            } else {
-                ui.label(ctx.re_ui.error_text("Error: missing instance key"));
-            }
+            ctx.component_ui_registry.ui(
+                ctx,
+                ui,
+                verbosity,
+                query,
+                db,
+                &self.entity_path,
+                &self.results,
+                &re_types::components::InstanceKey(0),
+            );
         } else if one_line {
             ui.label(format!("{} values", re_format::format_uint(num_instances)));
         } else {
@@ -106,49 +100,49 @@ impl DataUi for EntityComponentWithInstances {
                 .header(re_ui::ReUi::table_header_height(), |mut header| {
                     re_ui::ReUi::setup_table_header(&mut header);
                     header.col(|ui| {
-                        ui.label("Index");
+                        ui.label("Instance");
                     });
                     header.col(|ui| {
-                        ui.label(self.component_name().short_name());
+                        ui.label(self.component_name.short_name());
                     });
                 })
                 .body(|mut body| {
                     re_ui::ReUi::setup_table_body(&mut body);
                     let row_height = re_ui::ReUi::table_line_height();
-                    body.rows(row_height, displayed_row, |mut row| {
-                        if let Some(instance_key) = instance_keys.get(row.index()) {
-                            row.col(|ui| {
-                                let instance_path =
-                                    InstancePath::instance(self.entity_path.clone(), *instance_key);
-                                item_ui::instance_path_button_to(
-                                    ctx,
-                                    query,
-                                    store,
-                                    ui,
-                                    None,
-                                    &instance_path,
-                                    instance_key.syntax_highlighted(ui.style()),
-                                );
-                            });
-                            row.col(|ui| {
-                                ctx.component_ui_registry.ui(
-                                    ctx,
-                                    ui,
-                                    UiVerbosity::Small,
-                                    query,
-                                    store,
-                                    &self.entity_path,
-                                    &self.component_data,
-                                    instance_key,
-                                );
-                            });
-                        }
+                    body.rows(row_height, num_displayed_rows, |mut row| {
+                        let instance_key = re_types::components::InstanceKey(row.index() as _);
+                        row.col(|ui| {
+                            let instance_path =
+                                InstancePath::instance(self.entity_path.clone(), instance_key);
+                            item_ui::instance_path_button_to(
+                                ctx,
+                                query,
+                                db,
+                                ui,
+                                None,
+                                &instance_path,
+                                instance_key.syntax_highlighted(ui.style()),
+                            );
+                        });
+                        row.col(|ui| {
+                            ctx.component_ui_registry.ui(
+                                ctx,
+                                ui,
+                                UiVerbosity::Small,
+                                query,
+                                db,
+                                &self.entity_path,
+                                &self.results,
+                                &instance_key,
+                            );
+                        });
                     });
                 });
-            if num_instances > displayed_row {
+
+            if num_instances > num_displayed_rows {
                 ui.label(format!(
                     "…and {} more.",
-                    re_format::format_uint(num_instances - displayed_row)
+                    re_format::format_uint(num_instances - num_displayed_rows)
                 ));
             }
         }

crates/re_data_ui/src/component_path.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use re_log_types::ComponentPath;
 use re_viewer_context::{UiVerbosity, ViewerContext};
 
@@ -10,7 +12,7 @@ impl DataUi for ComponentPath {
         ui: &mut egui::Ui,
         verbosity: UiVerbosity,
         query: &re_data_store::LatestAtQuery,
-        store: &re_data_store::DataStore,
+        db: &re_entity_db::EntityDb,
     ) {
         let Self {
             entity_path,
@@ -21,27 +23,31 @@ impl DataUi for ComponentPath {
             ui.label(format!(
                 "Indicator component for the {archetype_name} archetype"
             ));
-        } else if let Some((_, _, component_data)) =
-            re_query::get_component_with_instances(store, query, entity_path, *component_name)
-        {
-            super::component::EntityComponentWithInstances {
-                entity_path: self.entity_path.clone(),
-                component_data,
-            }
-            .data_ui(ctx, ui, verbosity, query, store);
-        } else if let Some(entity_tree) = ctx.recording().tree().subtree(entity_path) {
-            if entity_tree.entity.components.contains_key(component_name) {
-                ui.label("<unset>");
+        } else {
+            let results =
+                db.query_caches2()
+                    .latest_at(db.store(), query, entity_path, [*component_name]);
+            if let Some(results) = results.components.get(component_name) {
+                crate::EntityLatestAtResults {
+                    entity_path: entity_path.clone(),
+                    component_name: *component_name,
+                    results: Arc::clone(results),
+                }
+                .data_ui(ctx, ui, verbosity, query, db);
+            } else if let Some(entity_tree) = ctx.recording().tree().subtree(entity_path) {
+                if entity_tree.entity.components.contains_key(component_name) {
+                    ui.label("<unset>");
+                } else {
+                    ui.label(format!(
+                        "Entity {entity_path:?} has no component {component_name:?}"
+                    ));
+                }
             } else {
-                ui.label(format!(
-                    "Entity {entity_path:?} has no component {component_name:?}"
-                ));
+                ui.label(
+                    ctx.re_ui
+                        .error_text(format!("Unknown component path: {self}")),
+                );
             }
-        } else {
-            ui.label(
-                ctx.re_ui
-                    .error_text(format!("Unknown component path: {self}")),
-            );
         }
     }
 }

crates/re_data_ui/src/component_ui_registry.rs

@@ -1,6 +1,6 @@
-use re_data_store::{DataStore, LatestAtQuery};
+use re_data_store::LatestAtQuery;
+use re_entity_db::{external::re_query_cache2::CachedLatestAtComponentResults, EntityDb};
 use re_log_types::{external::arrow2, EntityPath};
-use re_query::ComponentWithInstances;
 use re_types::external::arrow2::array::Utf8Array;
 use re_viewer_context::{ComponentUiRegistry, UiVerbosity, ViewerContext};
 
@@ -39,18 +39,13 @@ pub fn add_to_registry<C: EntityDataUi + re_types::Component>(registry: &mut Com
     registry.add(
         C::name(),
         Box::new(
-            |ctx, ui, verbosity, query, store, entity_path, component, instance| match component
-                .lookup::<C>(instance)
-            {
-                Ok(component) => {
-                    component.entity_data_ui(ctx, ui, verbosity, entity_path, query, store);
-                }
-                Err(re_query::QueryError::ComponentNotFound(_)) => {
+            |ctx, ui, verbosity, query, db, entity_path, component, instance| {
+                // TODO(#5607): what should happen if the promise is still pending?
+                if let Some(component) = component.instance::<C>(db.resolver(), instance.0 as _) {
+                    component.entity_data_ui(ctx, ui, verbosity, entity_path, query, db);
+                } else {
                     ui.weak("(not found)");
                 }
-                Err(err) => {
-                    re_log::warn_once!("Expected component {}, {}", C::name(), err);
-                }
             },
         ),
     );
@@ -62,13 +57,20 @@ fn fallback_component_ui(
     ui: &mut egui::Ui,
     verbosity: UiVerbosity,
     _query: &LatestAtQuery,
-    _store: &DataStore,
+    db: &EntityDb,
     _entity_path: &EntityPath,
-    component: &ComponentWithInstances,
+    component: &CachedLatestAtComponentResults,
     instance_key: &re_types::components::InstanceKey,
 ) {
+    // TODO(#5607): what should happen if the promise is still pending?
+    let value = component
+        .component_name(db.resolver())
+        .and_then(|component_name| {
+            component.instance_raw(db.resolver(), component_name, instance_key.0 as _)
+        });
+
     // No special ui implementation - use a generic one:
-    if let Some(value) = component.lookup_arrow(instance_key) {
+    if let Some(value) = value {
         arrow_ui(ui, verbosity, &*value);
     } else {
         ui.weak("(null)");