Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce DeserializableComponent trait and high-level query_latest #1417

Merged
merged 6 commits into from
Feb 27, 2023
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 15 additions & 12 deletions crates/re_data_store/src/entity_properties.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use re_arrow_store::LatestAtQuery;
use re_log_types::{
external::arrow2_convert::deserialize::arrow_array_deserialize_iterator, msg_bundle::Component,
EntityPath, Transform,
external::arrow2_convert::deserialize::arrow_array_deserialize_iterator,
msg_bundle::DeserializableComponent, EntityPath,
};

use crate::log_db::EntityDb;
Expand Down Expand Up @@ -38,6 +38,10 @@ pub struct EntityProperties {
pub visible: bool,
pub visible_history: ExtraQueryHistory,
pub interactive: bool,

/// Distance of the projection plane.
///
/// Only applies to pinhole cameras when in a spatial view, using 3D navigation.
pinhole_image_plane_distance: Option<ordered_float::NotNan<f32>>,
}

Expand Down Expand Up @@ -114,31 +118,30 @@ impl ExtraQueryHistory {

// ----------------------------------------------------------------------------

/// Get the latest value of the transform
///
/// We first look for the transform in the classic storage system since that's
/// what most users are still using. If we don't find the transform there, then
/// we check to see if it exists in the arrow storage.
pub fn query_transform(
/// Get the latest value for a given [`re_log_types::msg_bundle::Component`].
pub fn query_latest<C: DeserializableComponent>(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should there be something in the name or comment on how this expects that there is only a single component of the type on the entity?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's already guaranteed by the data model 🤔 Ha, you mean a single instance. Yes.

entity_db: &EntityDb,
entity_path: &EntityPath,
query: &LatestAtQuery,
) -> Option<Transform> {
) -> Option<C>
where
for<'b> &'b C::ArrayType: IntoIterator,
{
crate::profile_function!();

// Although it would be nice to use the `re_query` helpers for this, we would need to move
// this out of re_data_store to avoid a circular dep. Since we don't need to do a join for
// transforms this is easy enough.
let data_store = &entity_db.data_store;

let components = [Transform::name()];
let components = [C::name()];

let row_indices = data_store.latest_at(query, entity_path, Transform::name(), &components)?;
let row_indices = data_store.latest_at(query, entity_path, C::name(), &components)?;

let results = data_store.get(&components, &row_indices);
let arr = results.get(0)?.as_ref()?.as_ref();

let mut iter = arrow_array_deserialize_iterator::<Transform>(arr).ok()?;
let mut iter = arrow_array_deserialize_iterator::<C>(arr).ok()?;
teh-cmc marked this conversation as resolved.
Show resolved Hide resolved

let transform = iter.next();

Expand Down
33 changes: 30 additions & 3 deletions crates/re_log_types/src/msg_bundle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ use arrow2::{
offset::Offsets,
};
use arrow2_convert::{
deserialize::{ArrowArray, ArrowDeserialize},
field::ArrowField,
serialize::{ArrowSerialize, TryIntoArrow},
};
Expand Down Expand Up @@ -88,9 +89,8 @@ pub trait Component: ArrowField {
}
}

/// A trait to identify any [`Component`] that is ready to be collected and subsequently serialized
/// into an Arrow payload.
pub trait SerializableComponent
/// A [`Component`] that fulfils all the conditions required to be serialized as an Arrow payload.
pub trait SerializableComponent<ArrowFieldType = Self>
where
Self: Component + ArrowSerialize + ArrowField<Type = Self> + 'static,
{
Expand All @@ -101,6 +101,33 @@ impl<C> SerializableComponent for C where
{
}

/// A [`Component`] that fulfils all the conditions required to be deserialized from an Arrow
/// payload.
///
/// Note that due to the use of HRTBs in `arrow2_convert` traits, you will still need an extra HRTB
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HRTB for the uninitiated?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://doc.rust-lang.org/nomicon/hrtb.html

An old-school workaround for the lack of generic associated types... which actually have landed since then, at least in part.
We should be able to make those APIs in arrow2_convert friendlier pretty soon.

/// clause when marking a type as `DeserializableComponent`:
/// ```ignore
/// where
/// T: SerializableComponent,
/// for<'a> &'a T::ArrayType: IntoIterator,
/// ```
pub trait DeserializableComponent<ArrowFieldType = Self>
where
Self: Component + ArrowDeserialize + ArrowField<Type = ArrowFieldType> + 'static,
teh-cmc marked this conversation as resolved.
Show resolved Hide resolved
Self::ArrayType: ArrowArray,
for<'b> &'b Self::ArrayType: IntoIterator,
{
}

impl<C> DeserializableComponent for C
where
C: Component,
C: ArrowDeserialize + ArrowField<Type = C> + 'static,
C::ArrayType: ArrowArray,
for<'b> &'b C::ArrayType: IntoIterator,
{
}

/// A [`ComponentBundle`] holds an Arrow component column, and its field name.
///
/// A [`ComponentBundle`] can be created from a collection of any element that implements the
Expand Down
51 changes: 16 additions & 35 deletions crates/re_query/src/dataframe_util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,8 @@ use polars_core::prelude::*;
use re_arrow_store::ArrayExt;
use re_log_types::{
component_types::InstanceKey,
external::arrow2_convert::{
deserialize::{arrow_array_deserialize_iterator, ArrowArray, ArrowDeserialize},
field::ArrowField,
serialize::ArrowSerialize,
},
msg_bundle::Component,
external::arrow2_convert::deserialize::arrow_array_deserialize_iterator,
msg_bundle::{Component, DeserializableComponent, SerializableComponent},
};

use crate::{
Expand Down Expand Up @@ -55,10 +51,10 @@ fn fix_polars_nulls<C: Component>(array: &dyn Array) -> Box<dyn Array> {
}

/// Iterator for a single column in a dataframe as the rust-native Component type
pub fn iter_column<'a, C: Component>(df: &'a DataFrame) -> impl Iterator<Item = Option<C>> + 'a
pub fn iter_column<'a, C: DeserializableComponent>(
df: &'a DataFrame,
) -> impl Iterator<Item = Option<C>> + 'a
where
C: ArrowDeserialize + ArrowField<Type = C> + 'static,
C::ArrayType: ArrowArray,
for<'b> &'b C::ArrayType: IntoIterator,
{
let res = match df.column(C::name().as_str()) {
Expand All @@ -74,11 +70,7 @@ where
res.into_iter()
}

pub fn df_builder1<C0>(c0: &Vec<Option<C0>>) -> crate::Result<DataFrame>
where
C0: Component + 'static,
Option<C0>: ArrowSerialize + ArrowField<Type = Option<C0>>,
{
pub fn df_builder1<C0: SerializableComponent>(c0: &Vec<Option<C0>>) -> crate::Result<DataFrame> {
use arrow2::array::MutableArray;
use re_log_types::external::arrow2_convert::serialize::arrow_serialize_to_mutable_array;

Expand All @@ -92,10 +84,8 @@ where

pub fn df_builder2<C0, C1>(c0: &Vec<Option<C0>>, c1: &Vec<Option<C1>>) -> crate::Result<DataFrame>
where
C0: Component + 'static,
Option<C0>: ArrowSerialize + ArrowField<Type = Option<C0>>,
C1: Component + 'static,
Option<C1>: ArrowSerialize + ArrowField<Type = Option<C1>>,
C0: SerializableComponent,
C1: SerializableComponent,
{
use arrow2::array::MutableArray;
use re_log_types::external::arrow2_convert::serialize::arrow_serialize_to_mutable_array;
Expand All @@ -117,12 +107,9 @@ pub fn df_builder3<C0, C1, C2>(
c2: &Vec<Option<C2>>,
) -> crate::Result<DataFrame>
where
C0: Component + 'static,
Option<C0>: ArrowSerialize + ArrowField<Type = Option<C0>>,
C1: Component + 'static,
Option<C1>: ArrowSerialize + ArrowField<Type = Option<C1>>,
C2: Component + 'static,
Option<C2>: ArrowSerialize + ArrowField<Type = Option<C2>>,
C0: SerializableComponent,
C1: SerializableComponent,
C2: SerializableComponent,
{
use arrow2::array::MutableArray;
use re_log_types::external::arrow2_convert::serialize::arrow_serialize_to_mutable_array;
Expand All @@ -142,12 +129,10 @@ where
}

impl ComponentWithInstances {
pub fn as_df<C0>(&self) -> crate::Result<DataFrame>
pub fn as_df<C0: SerializableComponent + DeserializableComponent>(
&self,
) -> crate::Result<DataFrame>
where
C0: Component,
Option<C0>: ArrowSerialize + ArrowField<Type = Option<C0>>,
C0: ArrowDeserialize + ArrowField<Type = C0> + 'static,
C0::ArrayType: ArrowArray,
for<'a> &'a C0::ArrayType: IntoIterator,
{
if C0::name() != self.name {
Expand All @@ -169,8 +154,7 @@ impl ComponentWithInstances {

impl<Primary> EntityView<Primary>
where
Primary: Component + ArrowSerialize + ArrowDeserialize + ArrowField<Type = Primary> + 'static,
Primary::ArrayType: ArrowArray,
Primary: SerializableComponent + DeserializableComponent,
for<'a> &'a Primary::ArrayType: IntoIterator,
{
pub fn as_df1(&self) -> crate::Result<DataFrame> {
Expand All @@ -184,10 +168,7 @@ where

pub fn as_df2<C1>(&self) -> crate::Result<DataFrame>
where
C1: Clone + Component,
Option<C1>: ArrowSerialize + ArrowField<Type = Option<C1>>,
C1: ArrowDeserialize + ArrowField<Type = C1> + 'static,
C1::ArrayType: ArrowArray,
C1: SerializableComponent + DeserializableComponent + Clone,
for<'a> &'a C1::ArrayType: IntoIterator,
{
let instance_keys = self.primary.iter_instance_keys()?.map(Some).collect_vec();
Expand Down
32 changes: 10 additions & 22 deletions crates/re_query/src/entity_view.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,9 @@ use re_format::arrow;
use re_log_types::{
component_types::InstanceKey,
external::arrow2_convert::{
deserialize::{arrow_array_deserialize_iterator, ArrowArray, ArrowDeserialize},
field::ArrowField,
serialize::ArrowSerialize,
deserialize::arrow_array_deserialize_iterator, field::ArrowField, serialize::ArrowSerialize,
},
msg_bundle::Component,
msg_bundle::{Component, DeserializableComponent, SerializableComponent},
ComponentName,
};

Expand Down Expand Up @@ -58,10 +56,10 @@ impl ComponentWithInstances {
}

/// Iterate over the values and convert them to a native `Component`
pub fn iter_values<C: Component>(&self) -> crate::Result<impl Iterator<Item = Option<C>> + '_>
pub fn iter_values<C: DeserializableComponent>(
&self,
) -> crate::Result<impl Iterator<Item = Option<C>> + '_>
where
C: ArrowDeserialize + ArrowField<Type = C> + 'static,
C::ArrayType: ArrowArray,
for<'a> &'a C::ArrayType: IntoIterator,
{
if C::name() != self.name {
Expand All @@ -77,10 +75,8 @@ impl ComponentWithInstances {
}

/// Look up the value that corresponds to a given `InstanceKey` and convert to `Component`
pub fn lookup<C: Component>(&self, instance_key: &InstanceKey) -> crate::Result<C>
pub fn lookup<C: DeserializableComponent>(&self, instance_key: &InstanceKey) -> crate::Result<C>
where
C: ArrowDeserialize + ArrowField<Type = C> + 'static,
C::ArrayType: ArrowArray,
for<'a> &'a C::ArrayType: IntoIterator,
{
if C::name() != self.name {
Expand Down Expand Up @@ -132,14 +128,10 @@ impl ComponentWithInstances {
}

/// Produce a `ComponentWithInstances` from native component types
pub fn from_native<C>(
pub fn from_native<C: SerializableComponent>(
instance_keys: Option<&Vec<InstanceKey>>,
values: &Vec<C>,
) -> crate::Result<ComponentWithInstances>
where
C: Component + 'static,
C: ArrowSerialize + ArrowField<Type = C>,
{
) -> crate::Result<ComponentWithInstances> {
use re_log_types::external::arrow2_convert::serialize::arrow_serialize_to_mutable_array;

let instance_keys = if let Some(keys) = instance_keys {
Expand Down Expand Up @@ -296,10 +288,8 @@ where
}
}

impl<Primary> EntityView<Primary>
impl<Primary: SerializableComponent + DeserializableComponent> EntityView<Primary>
where
Primary: Component + ArrowSerialize + ArrowDeserialize + ArrowField<Type = Primary> + 'static,
Primary::ArrayType: ArrowArray,
for<'a> &'a Primary::ArrayType: IntoIterator,
{
/// Iterate over the instance keys
Expand Down Expand Up @@ -332,12 +322,10 @@ where
/// Iterate over the values of a `Component`.
///
/// Always produces an iterator of length `self.primary.len()`
pub fn iter_component<C: Component>(
pub fn iter_component<C: DeserializableComponent + Clone>(
&self,
) -> crate::Result<impl Iterator<Item = Option<C>> + '_>
where
C: Clone + ArrowDeserialize + ArrowField<Type = C> + 'static,
C::ArrayType: ArrowArray,
for<'b> &'b C::ArrayType: IntoIterator,
{
let component = self.components.get(&C::name());
Expand Down
14 changes: 3 additions & 11 deletions crates/re_query/src/visit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,7 @@

use re_log_types::{
component_types::InstanceKey,
external::arrow2_convert::{
deserialize::{ArrowArray, ArrowDeserialize},
field::ArrowField,
serialize::ArrowSerialize,
},
msg_bundle::Component,
msg_bundle::{Component, DeserializableComponent, SerializableComponent},
};

use crate::EntityView;
Expand All @@ -72,8 +67,7 @@ macro_rules! create_visitor {

) -> crate::Result<()>
where $(
$CC: Clone + ArrowDeserialize + ArrowField<Type = $CC> + 'static,
$CC::ArrayType: ArrowArray,
$CC: Clone + DeserializableComponent,
for<'a> &'a $CC::ArrayType: IntoIterator,
)*
{
Expand All @@ -98,10 +92,8 @@ macro_rules! create_visitor {
);
}

impl<Primary> EntityView<Primary>
impl<Primary: SerializableComponent + DeserializableComponent> EntityView<Primary>
where
Primary: Component + ArrowSerialize + ArrowDeserialize + ArrowField<Type = Primary> + 'static,
Primary::ArrayType: ArrowArray,
for<'a> &'a Primary::ArrayType: IntoIterator,
{
create_visitor! {visit1; ;}
Expand Down
6 changes: 3 additions & 3 deletions crates/re_viewer/src/misc/space_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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_transform, EntityPath, EntityTree};
use re_data_store::{log_db::EntityDb, query_latest, EntityPath, EntityTree};
use re_log_types::{Transform, ViewCoordinates};
use re_query::query_entity_with_primary;

Expand Down Expand Up @@ -112,7 +112,7 @@ impl SpaceInfoCollection {
tree: &EntityTree,
query: &LatestAtQuery,
) {
if let Some(transform) = query_transform(entity_db, &tree.path, query) {
if let Some(transform) = query_latest::<Transform>(entity_db, &tree.path, query) {
// A set transform (likely non-identity) - create a new space.
parent_space
.child_spaces
Expand Down Expand Up @@ -156,7 +156,7 @@ impl SpaceInfoCollection {
let mut spaces_info = Self::default();

// Start at the root. The root is always part of the collection!
if query_transform(entity_db, &EntityPath::root(), &query).is_some() {
if query_latest::<Transform>(entity_db, &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?");
}
let mut root_space_info = SpaceInfo::new(EntityPath::root());
Expand Down
6 changes: 3 additions & 3 deletions crates/re_viewer/src/misc/transform_cache.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use nohash_hasher::IntMap;
use re_arrow_store::LatestAtQuery;
use re_data_store::{log_db::EntityDb, query_transform, EntityPath, EntityPropertyMap, EntityTree};
use re_data_store::{log_db::EntityDb, query_latest, EntityPath, EntityPropertyMap, EntityTree};

use crate::misc::TimeControl;

Expand Down Expand Up @@ -222,7 +222,7 @@ fn transform_at(
query: &LatestAtQuery,
encountered_pinhole: &mut bool,
) -> Result<Option<macaw::Mat4>, UnreachableTransform> {
if let Some(transform) = query_transform(entity_db, entity_path, query) {
if let Some(transform) = query_latest(entity_db, entity_path, query) {
match transform {
re_log_types::Transform::Rigid3(rigid) => Ok(Some(rigid.parent_from_child().to_mat4())),
// If we're connected via 'unknown' it's not reachable
Expand Down Expand Up @@ -269,7 +269,7 @@ fn inverse_transform_at(
query: &LatestAtQuery,
encountered_pinhole: &mut bool,
) -> Result<Option<macaw::Mat4>, UnreachableTransform> {
if let Some(parent_transform) = query_transform(entity_db, entity_path, query) {
if let Some(parent_transform) = query_latest(entity_db, entity_path, query) {
match parent_transform {
re_log_types::Transform::Rigid3(rigid) => Ok(Some(rigid.child_from_parent().to_mat4())),
// If we're connected via 'unknown', everything except whats under `parent_tree` is unreachable
Expand Down
Loading