From 8c66f64023df457926c888ac05e9473b852d8ebc Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Thu, 21 Nov 2024 16:31:18 +0100 Subject: [PATCH] Add an arrow1-interfacr on top of Loggable --- Cargo.lock | 1 + crates/store/re_types_core/Cargo.toml | 2 + crates/store/re_types_core/src/loggable.rs | 64 ++++++++++++++++++++-- 3 files changed, 63 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2528ad7d681c..6a11c06acdda 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6342,6 +6342,7 @@ name = "re_types_core" version = "0.21.0-alpha.1+dev" dependencies = [ "anyhow", + "arrow", "backtrace", "bytemuck", "criterion", diff --git a/crates/store/re_types_core/Cargo.toml b/crates/store/re_types_core/Cargo.toml index ecc433d31a6c..5be97ce1e10e 100644 --- a/crates/store/re_types_core/Cargo.toml +++ b/crates/store/re_types_core/Cargo.toml @@ -42,7 +42,9 @@ re_tuid.workspace = true # External anyhow.workspace = true +arrow.workspace = true arrow2 = { workspace = true, features = [ + "arrow", "io_ipc", "io_print", "compute_concatenate", diff --git a/crates/store/re_types_core/src/loggable.rs b/crates/store/re_types_core/src/loggable.rs index 0c42840936d5..e331da70b718 100644 --- a/crates/store/re_types_core/src/loggable.rs +++ b/crates/store/re_types_core/src/loggable.rs @@ -16,6 +16,11 @@ use crate::{Archetype, ComponentBatch, LoggableBatch}; /// Implementing the [`Loggable`] trait automatically derives the [`LoggableBatch`] implementation, /// which makes it possible to work with lists' worth of data in a generic fashion. pub trait Loggable: 'static + Send + Sync + Clone + Sized + SizeBytes { + /// The underlying [`arrow::datatypes::DataType`], excluding datatype extensions. + fn arrow_datatype() -> arrow::datatypes::DataType { + Self::arrow2_datatype().into() + } + /// The underlying [`arrow2::datatypes::DataType`], excluding datatype extensions. fn arrow2_datatype() -> arrow2::datatypes::DataType; @@ -25,9 +30,24 @@ pub trait Loggable: 'static + Send + Sync + Clone + Sized + SizeBytes { /// When using Rerun's builtin components & datatypes, this can only fail if the data /// exceeds the maximum number of entries in an Arrow array (2^31 for standard arrays, /// 2^63 for large arrays). + fn to_arrow_opt<'a>( + data: impl IntoIterator>>>, + ) -> SerializationResult + where + Self: 'a, + { + Self::to_arrow2_opt(data).map(|array| arrow2::array::to_data(array.as_ref())) + } + + /// Given an iterator of options of owned or reference values to the current + /// [`Loggable`], serializes them into an Arrow2 array. + /// + /// When using Rerun's builtin components & datatypes, this can only fail if the data + /// exceeds the maximum number of entries in an Arrow2 array (2^31 for standard arrays, + /// 2^63 for large arrays). fn to_arrow2_opt<'a>( data: impl IntoIterator>>>, - ) -> SerializationResult> + ) -> SerializationResult> where Self: 'a; @@ -40,9 +60,26 @@ pub trait Loggable: 'static + Send + Sync + Clone + Sized + SizeBytes { /// exceeds the maximum number of entries in an Arrow array (2^31 for standard arrays, /// 2^63 for large arrays). #[inline] + fn to_arrow<'a>( + data: impl IntoIterator>>, + ) -> SerializationResult + where + Self: 'a, + { + re_tracing::profile_function!(); + Self::to_arrow_opt(data.into_iter().map(Some)) + } + + /// Given an iterator of owned or reference values to the current [`Loggable`], serializes + /// them into an Arrow2 array. + /// + /// When using Rerun's builtin components & datatypes, this can only fail if the data + /// exceeds the maximum number of entries in an Arrow2 array (2^31 for standard arrays, + /// 2^63 for large arrays). + #[inline] fn to_arrow2<'a>( data: impl IntoIterator>>, - ) -> SerializationResult> + ) -> SerializationResult> where Self: 'a, { @@ -54,7 +91,21 @@ pub trait Loggable: 'static + Send + Sync + Clone + Sized + SizeBytes { /// Given an Arrow array, deserializes it into a collection of [`Loggable`]s. #[inline] - fn from_arrow2(data: &dyn ::arrow2::array::Array) -> DeserializationResult> { + fn from_arrow(data: &arrow::array::ArrayData) -> DeserializationResult> { + re_tracing::profile_function!(); + Self::from_arrow_opt(data)? + .into_iter() + .map(|opt| { + opt.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: _Backtrace::new_unresolved(), + }) + }) + .collect::>>() + } + + /// Given an Arrow2 array, deserializes it into a collection of [`Loggable`]s. + #[inline] + fn from_arrow2(data: &dyn arrow2::array::Array) -> DeserializationResult> { re_tracing::profile_function!(); Self::from_arrow2_opt(data)? .into_iter() @@ -67,8 +118,13 @@ pub trait Loggable: 'static + Send + Sync + Clone + Sized + SizeBytes { } /// Given an Arrow array, deserializes it into a collection of optional [`Loggable`]s. + fn from_arrow_opt(data: &arrow::array::ArrayData) -> DeserializationResult>> { + Self::from_arrow2_opt(arrow2::array::from_data(data).as_ref()) + } + + /// Given an Arrow2 array, deserializes it into a collection of optional [`Loggable`]s. fn from_arrow2_opt( - data: &dyn ::arrow2::array::Array, + data: &dyn arrow2::array::Array, ) -> DeserializationResult>> { _ = data; // NOTE: do this here to avoid breaking users' autocomplete snippets Err(crate::DeserializationError::NotImplemented {