diff --git a/Cargo.lock b/Cargo.lock index 927ed6876abc..6251356e79dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4996,6 +4996,15 @@ dependencies = [ "rerun", ] +[[package]] +name = "roundtrip_points3d" +version = "0.8.0-alpha.6" +dependencies = [ + "anyhow", + "clap", + "rerun", +] + [[package]] name = "roundtrip_transform3d" version = "0.8.0-alpha.6" diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index f5856460f807..0127cc525a83 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -b4f2c2b3b6518fd5c5b8fe2255952c6986895efe58cbc6ae0b56b0384a9ad930 +53fa48636ea54f37bbe5b3b6c1e5479db495c6af3d95037c2b4b5a2ea0fc96a2 diff --git a/crates/re_types/src/archetypes/mod.rs b/crates/re_types/src/archetypes/mod.rs index b22ffe6bac10..255acf4f3435 100644 --- a/crates/re_types/src/archetypes/mod.rs +++ b/crates/re_types/src/archetypes/mod.rs @@ -2,8 +2,10 @@ mod fuzzy; mod points2d; +mod points3d; mod transform3d; pub use self::fuzzy::AffixFuzzer1; pub use self::points2d::Points2D; +pub use self::points3d::Points3D; pub use self::transform3d::Transform3D; diff --git a/crates/re_types/src/archetypes/points3d.rs b/crates/re_types/src/archetypes/points3d.rs new file mode 100644 index 000000000000..5afbc582c39c --- /dev/null +++ b/crates/re_types/src/archetypes/points3d.rs @@ -0,0 +1,597 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. + +#![allow(trivial_numeric_casts)] +#![allow(unused_parens)] +#![allow(clippy::clone_on_copy)] +#![allow(clippy::iter_on_single_items)] +#![allow(clippy::map_flatten)] +#![allow(clippy::match_wildcard_for_single_variants)] +#![allow(clippy::needless_question_mark)] +#![allow(clippy::redundant_closure)] +#![allow(clippy::too_many_arguments)] +#![allow(clippy::too_many_lines)] +#![allow(clippy::unnecessary_cast)] + +/// A 3D point cloud with positions and optional colors, radii, labels, etc. +/// +/// ## Example +/// +/// ```ignore +/// //! Log some very simple points. +/// use rerun::{experimental::archetypes::Points3D, MsgSender, RecordingStreamBuilder}; +/// +/// fn main() -> Result<(), Box> { +/// let (rec_stream, storage) = RecordingStreamBuilder::new("points").memory()?; +/// +/// MsgSender::from_archetype("points", &Points3D::new([(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]))? +/// .send(&rec_stream)?; +/// +/// rerun::native_viewer::show(storage.take())?; +/// Ok(()) +/// } +/// ``` +#[derive(Clone, Debug, PartialEq)] +pub struct Points3D { + /// All the actual 3D points that make up the point cloud. + pub points: Vec, + + /// Optional radii for the points, effectively turning them into circles. + pub radii: Option>, + + /// Optional colors for the points. + pub colors: Option>, + + /// Optional text labels for the points. + pub labels: Option>, + + /// An optional floating point value that specifies the 3D drawing order. + /// Objects with higher values are drawn on top of those with lower values. + /// + /// The default for 3D points is 30.0. + pub draw_order: Option, + + /// Optional class Ids for the points. + /// + /// The class ID provides colors and labels if not specified explicitly. + pub class_ids: Option>, + + /// Optional keypoint IDs for the points, identifying them within a class. + /// + /// If keypoint IDs are passed in but no class IDs were specified, the class ID will + /// default to 0. + /// This is useful to identify points within a single classification (which is identified + /// with `class_id`). + /// E.g. the classification might be 'Person' and the keypoints refer to joints on a + /// detected skeleton. + pub keypoint_ids: Option>, + + /// Unique identifiers for each individual point in the batch. + pub instance_keys: Option>, +} + +static REQUIRED_COMPONENTS: once_cell::sync::Lazy<[crate::ComponentName; 1usize]> = + once_cell::sync::Lazy::new(|| ["rerun.point3d".into()]); +static RECOMMENDED_COMPONENTS: once_cell::sync::Lazy<[crate::ComponentName; 2usize]> = + once_cell::sync::Lazy::new(|| ["rerun.radius".into(), "rerun.colorrgba".into()]); +static OPTIONAL_COMPONENTS: once_cell::sync::Lazy<[crate::ComponentName; 5usize]> = + once_cell::sync::Lazy::new(|| { + [ + "rerun.label".into(), + "rerun.draw_order".into(), + "rerun.class_id".into(), + "rerun.keypoint_id".into(), + "rerun.instance_key".into(), + ] + }); +static ALL_COMPONENTS: once_cell::sync::Lazy<[crate::ComponentName; 8usize]> = + once_cell::sync::Lazy::new(|| { + [ + "rerun.point3d".into(), + "rerun.radius".into(), + "rerun.colorrgba".into(), + "rerun.label".into(), + "rerun.draw_order".into(), + "rerun.class_id".into(), + "rerun.keypoint_id".into(), + "rerun.instance_key".into(), + ] + }); + +impl Points3D { + pub const NUM_COMPONENTS: usize = 8usize; +} + +impl crate::Archetype for Points3D { + #[inline] + fn name() -> crate::ArchetypeName { + crate::ArchetypeName::Borrowed("rerun.archetypes.Points3D") + } + + #[inline] + fn required_components() -> &'static [crate::ComponentName] { + REQUIRED_COMPONENTS.as_slice() + } + + #[inline] + fn recommended_components() -> &'static [crate::ComponentName] { + RECOMMENDED_COMPONENTS.as_slice() + } + + #[inline] + fn optional_components() -> &'static [crate::ComponentName] { + OPTIONAL_COMPONENTS.as_slice() + } + + #[inline] + fn all_components() -> &'static [crate::ComponentName] { + ALL_COMPONENTS.as_slice() + } + + #[inline] + fn try_to_arrow( + &self, + ) -> crate::SerializationResult< + Vec<(::arrow2::datatypes::Field, Box)>, + > { + use crate::Loggable as _; + Ok([ + { + Some({ + let array = + ::try_to_arrow(self.points.iter(), None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.Point3D".into(), + Box::new(array.data_type().clone()), + Some("rerun.point3d".into()), + ); + ( + ::arrow2::datatypes::Field::new("points", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#points".into(), + source: Box::new(err), + })? + }, + { + self.radii + .as_ref() + .map(|many| { + let array = ::try_to_arrow(many.iter(), None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.Radius".into(), + Box::new(array.data_type().clone()), + Some("rerun.radius".into()), + ); + ( + ::arrow2::datatypes::Field::new("radii", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#radii".into(), + source: Box::new(err), + })? + }, + { + self.colors + .as_ref() + .map(|many| { + let array = ::try_to_arrow(many.iter(), None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.Color".into(), + Box::new(array.data_type().clone()), + Some("rerun.colorrgba".into()), + ); + ( + ::arrow2::datatypes::Field::new("colors", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#colors".into(), + source: Box::new(err), + })? + }, + { + self.labels + .as_ref() + .map(|many| { + let array = ::try_to_arrow(many.iter(), None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.Label".into(), + Box::new(array.data_type().clone()), + Some("rerun.label".into()), + ); + ( + ::arrow2::datatypes::Field::new("labels", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#labels".into(), + source: Box::new(err), + })? + }, + { + self.draw_order + .as_ref() + .map(|single| { + let array = ::try_to_arrow([single], None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.DrawOrder".into(), + Box::new(array.data_type().clone()), + Some("rerun.draw_order".into()), + ); + ( + ::arrow2::datatypes::Field::new("draw_order", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#draw_order".into(), + source: Box::new(err), + })? + }, + { + self.class_ids + .as_ref() + .map(|many| { + let array = ::try_to_arrow(many.iter(), None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.ClassId".into(), + Box::new(array.data_type().clone()), + Some("rerun.class_id".into()), + ); + ( + ::arrow2::datatypes::Field::new("class_ids", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#class_ids".into(), + source: Box::new(err), + })? + }, + { + self.keypoint_ids + .as_ref() + .map(|many| { + let array = + ::try_to_arrow(many.iter(), None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.KeypointId".into(), + Box::new(array.data_type().clone()), + Some("rerun.keypoint_id".into()), + ); + ( + ::arrow2::datatypes::Field::new("keypoint_ids", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#keypoint_ids".into(), + source: Box::new(err), + })? + }, + { + self.instance_keys + .as_ref() + .map(|many| { + let array = + ::try_to_arrow(many.iter(), None); + array.map(|array| { + let datatype = ::arrow2::datatypes::DataType::Extension( + "rerun.components.InstanceKey".into(), + Box::new(array.data_type().clone()), + Some("rerun.instance_key".into()), + ); + ( + ::arrow2::datatypes::Field::new("instance_keys", datatype, false), + array, + ) + }) + }) + .transpose() + .map_err(|err| crate::SerializationError::Context { + location: "rerun.archetypes.Points3D#instance_keys".into(), + source: Box::new(err), + })? + }, + ] + .into_iter() + .flatten() + .collect()) + } + + #[inline] + fn try_from_arrow( + data: impl IntoIterator)>, + ) -> crate::DeserializationResult { + use crate::Loggable as _; + let arrays_by_name: ::std::collections::HashMap<_, _> = data + .into_iter() + .map(|(field, array)| (field.name, array)) + .collect(); + let points = { + let array = arrays_by_name + .get("points") + .ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#points".into(), + source: Box::new(err), + })?; + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#points".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#points".into(), + source: Box::new(err), + })? + }; + let radii = if let Some(array) = arrays_by_name.get("radii") { + Some( + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#radii".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#radii".into(), + source: Box::new(err), + })?, + ) + } else { + None + }; + let colors = if let Some(array) = arrays_by_name.get("colors") { + Some( + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#colors".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#colors".into(), + source: Box::new(err), + })?, + ) + } else { + None + }; + let labels = if let Some(array) = arrays_by_name.get("labels") { + Some( + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#labels".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#labels".into(), + source: Box::new(err), + })?, + ) + } else { + None + }; + let draw_order = if let Some(array) = arrays_by_name.get("draw_order") { + Some( + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#draw_order".into(), + source: Box::new(err), + })? + .into_iter() + .next() + .flatten() + .ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#draw_order".into(), + source: Box::new(err), + })?, + ) + } else { + None + }; + let class_ids = if let Some(array) = arrays_by_name.get("class_ids") { + Some( + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#class_ids".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#class_ids".into(), + source: Box::new(err), + })?, + ) + } else { + None + }; + let keypoint_ids = if let Some(array) = arrays_by_name.get("keypoint_ids") { + Some( + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#keypoint_ids".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#keypoint_ids".into(), + source: Box::new(err), + })?, + ) + } else { + None + }; + let instance_keys = if let Some(array) = arrays_by_name.get("instance_keys") { + Some( + ::try_from_arrow_opt(&**array) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#instance_keys".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.archetypes.Points3D#instance_keys".into(), + source: Box::new(err), + })?, + ) + } else { + None + }; + Ok(Self { + points, + radii, + colors, + labels, + draw_order, + class_ids, + keypoint_ids, + instance_keys, + }) + } +} + +impl Points3D { + pub fn new(points: impl IntoIterator>) -> Self { + Self { + points: points.into_iter().map(Into::into).collect(), + radii: None, + colors: None, + labels: None, + draw_order: None, + class_ids: None, + keypoint_ids: None, + instance_keys: None, + } + } + + pub fn with_radii( + mut self, + radii: impl IntoIterator>, + ) -> Self { + self.radii = Some(radii.into_iter().map(Into::into).collect()); + self + } + + pub fn with_colors( + mut self, + colors: impl IntoIterator>, + ) -> Self { + self.colors = Some(colors.into_iter().map(Into::into).collect()); + self + } + + pub fn with_labels( + mut self, + labels: impl IntoIterator>, + ) -> Self { + self.labels = Some(labels.into_iter().map(Into::into).collect()); + self + } + + pub fn with_draw_order(mut self, draw_order: impl Into) -> Self { + self.draw_order = Some(draw_order.into()); + self + } + + pub fn with_class_ids( + mut self, + class_ids: impl IntoIterator>, + ) -> Self { + self.class_ids = Some(class_ids.into_iter().map(Into::into).collect()); + self + } + + pub fn with_keypoint_ids( + mut self, + keypoint_ids: impl IntoIterator>, + ) -> Self { + self.keypoint_ids = Some(keypoint_ids.into_iter().map(Into::into).collect()); + self + } + + pub fn with_instance_keys( + mut self, + instance_keys: impl IntoIterator>, + ) -> Self { + self.instance_keys = Some(instance_keys.into_iter().map(Into::into).collect()); + self + } +} diff --git a/crates/re_types/src/components/mod.rs b/crates/re_types/src/components/mod.rs index 9a39229240d3..6669a0aead4b 100644 --- a/crates/re_types/src/components/mod.rs +++ b/crates/re_types/src/components/mod.rs @@ -15,6 +15,8 @@ mod label; mod label_ext; mod point2d; mod point2d_ext; +mod point3d; +mod point3d_ext; mod radius; mod radius_ext; mod transform3d; @@ -33,5 +35,6 @@ pub use self::instance_key::InstanceKey; pub use self::keypoint_id::KeypointId; pub use self::label::Label; pub use self::point2d::Point2D; +pub use self::point3d::Point3D; pub use self::radius::Radius; pub use self::transform3d::Transform3D; diff --git a/crates/re_types/src/components/point3d.rs b/crates/re_types/src/components/point3d.rs new file mode 100644 index 000000000000..23f585590808 --- /dev/null +++ b/crates/re_types/src/components/point3d.rs @@ -0,0 +1,138 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. + +#![allow(trivial_numeric_casts)] +#![allow(unused_parens)] +#![allow(clippy::clone_on_copy)] +#![allow(clippy::iter_on_single_items)] +#![allow(clippy::map_flatten)] +#![allow(clippy::match_wildcard_for_single_variants)] +#![allow(clippy::needless_question_mark)] +#![allow(clippy::redundant_closure)] +#![allow(clippy::too_many_arguments)] +#![allow(clippy::too_many_lines)] +#![allow(clippy::unnecessary_cast)] + +/// A point in 3D space. +#[derive(Clone, Debug, Default, Copy, PartialEq, PartialOrd)] +pub struct Point3D(pub crate::datatypes::Point3D); + +impl From for Point3D { + fn from(v: crate::datatypes::Point3D) -> Self { + Self(v) + } +} + +impl<'a> From for ::std::borrow::Cow<'a, Point3D> { + #[inline] + fn from(value: Point3D) -> Self { + std::borrow::Cow::Owned(value) + } +} + +impl<'a> From<&'a Point3D> for ::std::borrow::Cow<'a, Point3D> { + #[inline] + fn from(value: &'a Point3D) -> Self { + std::borrow::Cow::Borrowed(value) + } +} + +impl crate::Loggable for Point3D { + type Name = crate::ComponentName; + #[inline] + fn name() -> Self::Name { + "rerun.point3d".into() + } + + #[allow(unused_imports, clippy::wildcard_imports)] + #[inline] + fn to_arrow_datatype() -> arrow2::datatypes::DataType { + use ::arrow2::datatypes::*; + DataType::Struct(vec![ + Field { + name: "x".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "y".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "z".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + ]) + } + + #[allow(unused_imports, clippy::wildcard_imports)] + fn try_to_arrow_opt<'a>( + data: impl IntoIterator>>>, + extension_wrapper: Option<&str>, + ) -> crate::SerializationResult> + where + Self: Clone + 'a, + { + use crate::Loggable as _; + use ::arrow2::{array::*, datatypes::*}; + Ok({ + let (somes, data0): (Vec<_>, Vec<_>) = data + .into_iter() + .map(|datum| { + let datum: Option<::std::borrow::Cow<'a, Self>> = datum.map(Into::into); + let datum = datum.map(|datum| { + let Self(data0) = datum.into_owned(); + data0 + }); + (datum.is_some(), datum) + }) + .unzip(); + let data0_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + { + _ = data0_bitmap; + _ = extension_wrapper; + crate::datatypes::Point3D::try_to_arrow_opt( + data0, + Some("rerun.components.Point3D"), + )? + } + }) + } + + #[allow(unused_imports, clippy::wildcard_imports)] + fn try_from_arrow_opt( + data: &dyn ::arrow2::array::Array, + ) -> crate::DeserializationResult>> + where + Self: Sized, + { + use crate::Loggable as _; + use ::arrow2::{array::*, datatypes::*}; + Ok(crate::datatypes::Point3D::try_from_arrow_opt(data) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.components.Point3D#xy".into(), + source: Box::new(err), + })? + .into_iter() + .map(|v| { + v.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + }) + .map(|res| res.map(|v| Some(Self(v)))) + .collect::>>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.components.Point3D#xy".into(), + source: Box::new(err), + })?) + } +} + +impl crate::Component for Point3D {} diff --git a/crates/re_types/src/datatypes/mod.rs b/crates/re_types/src/datatypes/mod.rs index 2229fae46c03..157cf6f81faf 100644 --- a/crates/re_types/src/datatypes/mod.rs +++ b/crates/re_types/src/datatypes/mod.rs @@ -8,6 +8,8 @@ mod mat3x3_ext; mod mat4x4; mod point2d; mod point2d_ext; +mod point3d; +mod point3d_ext; mod quaternion; mod quaternion_ext; mod rotation3d; @@ -35,6 +37,7 @@ pub use self::fuzzy::{ pub use self::mat3x3::Mat3x3; pub use self::mat4x4::Mat4x4; pub use self::point2d::Point2D; +pub use self::point3d::Point3D; pub use self::quaternion::Quaternion; pub use self::rotation3d::Rotation3D; pub use self::rotation_axis_angle::RotationAxisAngle; diff --git a/crates/re_types/src/datatypes/point3d.rs b/crates/re_types/src/datatypes/point3d.rs new file mode 100644 index 000000000000..5b25b457ec6d --- /dev/null +++ b/crates/re_types/src/datatypes/point3d.rs @@ -0,0 +1,294 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. + +#![allow(trivial_numeric_casts)] +#![allow(unused_parens)] +#![allow(clippy::clone_on_copy)] +#![allow(clippy::iter_on_single_items)] +#![allow(clippy::map_flatten)] +#![allow(clippy::match_wildcard_for_single_variants)] +#![allow(clippy::needless_question_mark)] +#![allow(clippy::redundant_closure)] +#![allow(clippy::too_many_arguments)] +#![allow(clippy::too_many_lines)] +#![allow(clippy::unnecessary_cast)] + +/// A point in 3D space. +#[derive(Clone, Debug, Default, Copy, PartialEq, PartialOrd)] +pub struct Point3D { + pub x: f32, + pub y: f32, + pub z: f32, +} + +impl<'a> From for ::std::borrow::Cow<'a, Point3D> { + #[inline] + fn from(value: Point3D) -> Self { + std::borrow::Cow::Owned(value) + } +} + +impl<'a> From<&'a Point3D> for ::std::borrow::Cow<'a, Point3D> { + #[inline] + fn from(value: &'a Point3D) -> Self { + std::borrow::Cow::Borrowed(value) + } +} + +impl crate::Loggable for Point3D { + type Name = crate::DatatypeName; + #[inline] + fn name() -> Self::Name { + "rerun.datatypes.Point3D".into() + } + + #[allow(unused_imports, clippy::wildcard_imports)] + #[inline] + fn to_arrow_datatype() -> arrow2::datatypes::DataType { + use ::arrow2::datatypes::*; + DataType::Struct(vec![ + Field { + name: "x".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "y".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + Field { + name: "z".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }, + ]) + } + + #[allow(unused_imports, clippy::wildcard_imports)] + fn try_to_arrow_opt<'a>( + data: impl IntoIterator>>>, + extension_wrapper: Option<&str>, + ) -> crate::SerializationResult> + where + Self: Clone + 'a, + { + use crate::Loggable as _; + use ::arrow2::{array::*, datatypes::*}; + Ok({ + let (somes, data): (Vec<_>, Vec<_>) = data + .into_iter() + .map(|datum| { + let datum: Option<::std::borrow::Cow<'a, Self>> = datum.map(Into::into); + (datum.is_some(), datum) + }) + .unzip(); + let bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + StructArray::new( + (if let Some(ext) = extension_wrapper { + DataType::Extension( + ext.to_owned(), + Box::new(::to_arrow_datatype()), + None, + ) + } else { + ::to_arrow_datatype() + }) + .to_logical_type() + .clone(), + vec![ + { + let (somes, x): (Vec<_>, Vec<_>) = data + .iter() + .map(|datum| { + let datum = datum.as_ref().map(|datum| { + let Self { x, .. } = &**datum; + x.clone() + }); + (datum.is_some(), datum) + }) + .unzip(); + let x_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + PrimitiveArray::new( + { + _ = extension_wrapper; + DataType::Float32.to_logical_type().clone() + }, + x.into_iter().map(|v| v.unwrap_or_default()).collect(), + x_bitmap, + ) + .boxed() + }, + { + let (somes, y): (Vec<_>, Vec<_>) = data + .iter() + .map(|datum| { + let datum = datum.as_ref().map(|datum| { + let Self { y, .. } = &**datum; + y.clone() + }); + (datum.is_some(), datum) + }) + .unzip(); + let y_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + PrimitiveArray::new( + { + _ = extension_wrapper; + DataType::Float32.to_logical_type().clone() + }, + y.into_iter().map(|v| v.unwrap_or_default()).collect(), + y_bitmap, + ) + .boxed() + }, + { + let (somes, z): (Vec<_>, Vec<_>) = data + .iter() + .map(|datum| { + let datum = datum.as_ref().map(|datum| { + let Self { z, .. } = &**datum; + z.clone() + }); + (datum.is_some(), datum) + }) + .unzip(); + let z_bitmap: Option<::arrow2::bitmap::Bitmap> = { + let any_nones = somes.iter().any(|some| !*some); + any_nones.then(|| somes.into()) + }; + PrimitiveArray::new( + { + _ = extension_wrapper; + DataType::Float32.to_logical_type().clone() + }, + z.into_iter().map(|v| v.unwrap_or_default()).collect(), + z_bitmap, + ) + .boxed() + }, + ], + bitmap, + ) + .boxed() + }) + } + + #[allow(unused_imports, clippy::wildcard_imports)] + fn try_from_arrow_opt( + data: &dyn ::arrow2::array::Array, + ) -> crate::DeserializationResult>> + where + Self: Sized, + { + use crate::Loggable as _; + use ::arrow2::{array::*, datatypes::*}; + Ok({ + let data = data + .as_any() + .downcast_ref::<::arrow2::array::StructArray>() + .ok_or_else(|| crate::DeserializationError::DatatypeMismatch { + expected: data.data_type().clone(), + got: data.data_type().clone(), + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.datatypes.Point3D".into(), + source: Box::new(err), + })?; + if data.is_empty() { + Vec::new() + } else { + let (data_fields, data_arrays, data_bitmap) = + (data.fields(), data.values(), data.validity()); + let is_valid = |i| data_bitmap.map_or(true, |bitmap| bitmap.get_bit(i)); + let arrays_by_name: ::std::collections::HashMap<_, _> = data_fields + .iter() + .map(|field| field.name.as_str()) + .zip(data_arrays) + .collect(); + let x = { + let data = &**arrays_by_name["x"]; + + data.as_any() + .downcast_ref::() + .unwrap() + .into_iter() + .map(|v| v.copied()) + }; + let y = { + let data = &**arrays_by_name["y"]; + + data.as_any() + .downcast_ref::() + .unwrap() + .into_iter() + .map(|v| v.copied()) + }; + let z = { + let data = &**arrays_by_name["z"]; + + data.as_any() + .downcast_ref::() + .unwrap() + .into_iter() + .map(|v| v.copied()) + }; + ::itertools::izip!(x, y, z) + .enumerate() + .map(|(i, (x, y, z))| { + is_valid(i) + .then(|| { + Ok(Self { + x: x.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + .map_err(|err| { + crate::DeserializationError::Context { + location: "rerun.datatypes.Point3D#x".into(), + source: Box::new(err), + } + })?, + y: y.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + .map_err(|err| { + crate::DeserializationError::Context { + location: "rerun.datatypes.Point3D#y".into(), + source: Box::new(err), + } + })?, + z: z.ok_or_else(|| crate::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + .map_err(|err| { + crate::DeserializationError::Context { + location: "rerun.datatypes.Point3D#z".into(), + source: Box::new(err), + } + })?, + }) + }) + .transpose() + }) + .collect::>>() + .map_err(|err| crate::DeserializationError::Context { + location: "rerun.datatypes.Point3D".into(), + source: Box::new(err), + })? + } + }) + } +} + +impl crate::Datatype for Point3D {} diff --git a/rerun_cpp/src/archetypes.hpp b/rerun_cpp/src/archetypes.hpp index e69cd7d5414f..e1f509fd7d78 100644 --- a/rerun_cpp/src/archetypes.hpp +++ b/rerun_cpp/src/archetypes.hpp @@ -4,4 +4,5 @@ #include "archetypes/affix_fuzzer1.hpp" #include "archetypes/points2d.hpp" +#include "archetypes/points3d.hpp" #include "archetypes/transform3d.hpp" diff --git a/rerun_cpp/src/archetypes/points3d.cpp b/rerun_cpp/src/archetypes/points3d.cpp new file mode 100644 index 000000000000..c1830b732e53 --- /dev/null +++ b/rerun_cpp/src/archetypes/points3d.cpp @@ -0,0 +1,8 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. +// Based on "crates/re_types/definitions/rerun/archetypes/points3d.fbs" + +#include "points3d.hpp" + +namespace rr { + namespace archetypes {} +} // namespace rr diff --git a/rerun_cpp/src/archetypes/points3d.hpp b/rerun_cpp/src/archetypes/points3d.hpp new file mode 100644 index 000000000000..f483b05f1cbd --- /dev/null +++ b/rerun_cpp/src/archetypes/points3d.hpp @@ -0,0 +1,59 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. +// Based on "crates/re_types/definitions/rerun/archetypes/points3d.fbs" + +#pragma once + +#include +#include +#include + +#include "../components/class_id.hpp" +#include "../components/color.hpp" +#include "../components/draw_order.hpp" +#include "../components/instance_key.hpp" +#include "../components/keypoint_id.hpp" +#include "../components/label.hpp" +#include "../components/point3d.hpp" +#include "../components/radius.hpp" + +namespace rr { + namespace archetypes { + /// A 3D point cloud with positions and optional colors, radii, labels, etc. + struct Points3D { + /// All the actual 3D points that make up the point cloud. + std::vector points; + + /// Optional radii for the points, effectively turning them into circles. + std::optional> radii; + + /// Optional colors for the points. + std::optional> colors; + + /// Optional text labels for the points. + std::optional> labels; + + /// An optional floating point value that specifies the 3D drawing order. + /// Objects with higher values are drawn on top of those with lower values. + /// + /// The default for 3D points is 30.0. + std::optional draw_order; + + /// Optional class Ids for the points. + /// + /// The class ID provides colors and labels if not specified explicitly. + std::optional> class_ids; + + /// Optional keypoint IDs for the points, identifying them within a class. + /// + /// If keypoint IDs are passed in but no class IDs were specified, the class ID will + /// default to 0. + /// This is useful to identify points within a single classification (which is + /// identified with `class_id`). E.g. the classification might be 'Person' and the + /// keypoints refer to joints on a detected skeleton. + std::optional> keypoint_ids; + + /// Unique identifiers for each individual point in the batch. + std::optional> instance_keys; + }; + } // namespace archetypes +} // namespace rr diff --git a/rerun_cpp/src/components.hpp b/rerun_cpp/src/components.hpp index d72a4e6bbe63..5ce763a6c17a 100644 --- a/rerun_cpp/src/components.hpp +++ b/rerun_cpp/src/components.hpp @@ -28,5 +28,6 @@ #include "components/keypoint_id.hpp" #include "components/label.hpp" #include "components/point2d.hpp" +#include "components/point3d.hpp" #include "components/radius.hpp" #include "components/transform3d.hpp" diff --git a/rerun_cpp/src/components/point3d.cpp b/rerun_cpp/src/components/point3d.cpp new file mode 100644 index 000000000000..e1c46f971df0 --- /dev/null +++ b/rerun_cpp/src/components/point3d.cpp @@ -0,0 +1,15 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. +// Based on "crates/re_types/definitions/rerun/components/point3d.fbs" + +#include + +#include "../datatypes/point3d.hpp" +#include "point3d.hpp" + +namespace rr { + namespace components { + std::shared_ptr Point3D::to_arrow_datatype() { + return rr::datatypes::Point3D::to_arrow_datatype(); + } + } // namespace components +} // namespace rr diff --git a/rerun_cpp/src/components/point3d.hpp b/rerun_cpp/src/components/point3d.hpp new file mode 100644 index 000000000000..16341551009b --- /dev/null +++ b/rerun_cpp/src/components/point3d.hpp @@ -0,0 +1,29 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. +// Based on "crates/re_types/definitions/rerun/components/point3d.fbs" + +#pragma once + +#include +#include +#include + +#include "../datatypes/point3d.hpp" + +namespace arrow { + class DataType; +} + +namespace rr { + namespace components { + /// A point in 3D space. + struct Point3D { + rr::datatypes::Point3D xy; + + public: + Point3D(rr::datatypes::Point3D xy) : xy(std::move(xy)) {} + + /// Returns the arrow data type this type corresponds to. + static std::shared_ptr to_arrow_datatype(); + }; + } // namespace components +} // namespace rr diff --git a/rerun_cpp/src/datatypes.hpp b/rerun_cpp/src/datatypes.hpp index 65bfaa796e8e..d83dbcdbef83 100644 --- a/rerun_cpp/src/datatypes.hpp +++ b/rerun_cpp/src/datatypes.hpp @@ -12,6 +12,7 @@ #include "datatypes/mat3x3.hpp" #include "datatypes/mat4x4.hpp" #include "datatypes/point2d.hpp" +#include "datatypes/point3d.hpp" #include "datatypes/quaternion.hpp" #include "datatypes/rotation3d.hpp" #include "datatypes/rotation_axis_angle.hpp" diff --git a/rerun_cpp/src/datatypes/point3d.cpp b/rerun_cpp/src/datatypes/point3d.cpp new file mode 100644 index 000000000000..52e1c58d0aa0 --- /dev/null +++ b/rerun_cpp/src/datatypes/point3d.cpp @@ -0,0 +1,18 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. +// Based on "crates/re_types/definitions/rerun/datatypes/point3d.fbs" + +#include + +#include "point3d.hpp" + +namespace rr { + namespace datatypes { + std::shared_ptr Point3D::to_arrow_datatype() { + return arrow::struct_({ + arrow::field("x", arrow::float32(), false, nullptr), + arrow::field("y", arrow::float32(), false, nullptr), + arrow::field("z", arrow::float32(), false, nullptr), + }); + } + } // namespace datatypes +} // namespace rr diff --git a/rerun_cpp/src/datatypes/point3d.hpp b/rerun_cpp/src/datatypes/point3d.hpp new file mode 100644 index 000000000000..78b6e9b04e0b --- /dev/null +++ b/rerun_cpp/src/datatypes/point3d.hpp @@ -0,0 +1,28 @@ +// NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. +// Based on "crates/re_types/definitions/rerun/datatypes/point3d.fbs" + +#pragma once + +#include +#include + +namespace arrow { + class DataType; +} + +namespace rr { + namespace datatypes { + /// A point in 3D space. + struct Point3D { + float x; + + float y; + + float z; + + public: + /// Returns the arrow data type this type corresponds to. + static std::shared_ptr to_arrow_datatype(); + }; + } // namespace datatypes +} // namespace rr diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/__init__.py index d33d6f1f9829..2b77ae243ed7 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/__init__.py @@ -2,6 +2,6 @@ from __future__ import annotations -__all__ = ["AffixFuzzer1", "Points2D", "Transform3D"] +__all__ = ["AffixFuzzer1", "Points2D", "Points3D", "Transform3D"] -from .archetypes import AffixFuzzer1, Points2D, Transform3D +from .archetypes import AffixFuzzer1, Points2D, Points3D, Transform3D diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py index 70720d6c332e..a70cccfadd1e 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py @@ -4,6 +4,7 @@ from .fuzzy import AffixFuzzer1 from .points2d import Points2D +from .points3d import Points3D from .transform3d import Transform3D -__all__ = ["AffixFuzzer1", "Points2D", "Transform3D"] +__all__ = ["AffixFuzzer1", "Points2D", "Points3D", "Transform3D"] diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points3d.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points3d.py new file mode 100644 index 000000000000..e7bb53df9eb8 --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points3d.py @@ -0,0 +1,119 @@ +# NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. + +from __future__ import annotations + +from attrs import define, field + +from .. import components +from .._baseclasses import ( + Archetype, +) + +__all__ = ["Points3D"] + + +@define(str=False, repr=False) +class Points3D(Archetype): + """ + A 3D point cloud with positions and optional colors, radii, labels, etc. + + Example + ------- + ```python + import rerun as rr + import rerun.experimental as rr_exp + + rr.init("points", spawn=True) + + rr_exp.log_any("simple", rr_exp.Points3D([[0, 0, 0], [1, 1, 1]])) + ``` + """ + + points: components.Point3DArray = field( + metadata={"component": "primary"}, + converter=components.Point3DArray.from_similar, # type: ignore[misc] + ) + """ + All the actual 3D points that make up the point cloud. + """ + + radii: components.RadiusArray | None = field( + metadata={"component": "secondary"}, + default=None, + converter=components.RadiusArray.from_similar, # type: ignore[misc] + ) + """ + Optional radii for the points, effectively turning them into circles. + """ + + colors: components.ColorArray | None = field( + metadata={"component": "secondary"}, + default=None, + converter=components.ColorArray.from_similar, # type: ignore[misc] + ) + """ + Optional colors for the points. + + The colors are interpreted as RGB or RGBA in sRGB gamma-space, + As either 0-1 floats or 0-255 integers, with separate alpha. + """ + + labels: components.LabelArray | None = field( + metadata={"component": "secondary"}, + default=None, + converter=components.LabelArray.from_similar, # type: ignore[misc] + ) + """ + Optional text labels for the points. + """ + + draw_order: components.DrawOrderArray | None = field( + metadata={"component": "secondary"}, + default=None, + converter=components.DrawOrderArray.from_similar, # type: ignore[misc] + ) + """ + An optional floating point value that specifies the 3D drawing order. + Objects with higher values are drawn on top of those with lower values. + + The default for 3D points is 30.0. + """ + + class_ids: components.ClassIdArray | None = field( + metadata={"component": "secondary"}, + default=None, + converter=components.ClassIdArray.from_similar, # type: ignore[misc] + ) + """ + Optional class Ids for the points. + + The class ID provides colors and labels if not specified explicitly. + """ + + keypoint_ids: components.KeypointIdArray | None = field( + metadata={"component": "secondary"}, + default=None, + converter=components.KeypointIdArray.from_similar, # type: ignore[misc] + ) + """ + Optional keypoint IDs for the points, identifying them within a class. + + If keypoint IDs are passed in but no class IDs were specified, the class ID will + default to 0. + This is useful to identify points within a single classification (which is identified + with `class_id`). + E.g. the classification might be 'Person' and the keypoints refer to joints on a + detected skeleton. + """ + + instance_keys: components.InstanceKeyArray | None = field( + metadata={"component": "secondary"}, + default=None, + converter=components.InstanceKeyArray.from_similar, # type: ignore[misc] + ) + """ + Unique identifiers for each individual point in the batch. + """ + + __str__ = Archetype.__str__ + __repr__ = Archetype.__repr__ diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py index bf9870054250..62c5f65d3945 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py @@ -79,6 +79,7 @@ from .keypoint_id import KeypointId, KeypointIdArray, KeypointIdArrayLike, KeypointIdLike, KeypointIdType from .label import Label, LabelArray, LabelArrayLike, LabelLike, LabelType from .point2d import Point2DArray, Point2DType +from .point3d import Point3DArray, Point3DType from .radius import Radius, RadiusArray, RadiusArrayLike, RadiusLike, RadiusType from .transform3d import Transform3DArray, Transform3DType @@ -183,6 +184,8 @@ "LabelType", "Point2DArray", "Point2DType", + "Point3DArray", + "Point3DType", "Radius", "RadiusArray", "RadiusArrayLike", diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/point3d.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/point3d.py new file mode 100644 index 000000000000..4b63e3cced9f --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/point3d.py @@ -0,0 +1,28 @@ +# NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. + +from __future__ import annotations + +from .. import datatypes +from .._baseclasses import ( + BaseDelegatingExtensionArray, + BaseDelegatingExtensionType, +) + +__all__ = ["Point3DArray", "Point3DType"] + + +class Point3DType(BaseDelegatingExtensionType): + _TYPE_NAME = "rerun.point3d" + _DELEGATED_EXTENSION_TYPE = datatypes.Point3DType + + +class Point3DArray(BaseDelegatingExtensionArray[datatypes.Point3DArrayLike]): + _EXTENSION_NAME = "rerun.point3d" + _EXTENSION_TYPE = Point3DType + _DELEGATED_ARRAY_TYPE = datatypes.Point3DArray + + +Point3DType._ARRAY_TYPE = Point3DArray + +# TODO(cmc): bring back registration to pyarrow once legacy types are gone +# pa.register_extension_type(Point3DType()) diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py index 10f998ed4340..10a7a8f056b0 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py @@ -38,6 +38,7 @@ from .mat3x3 import Mat3x3, Mat3x3Array, Mat3x3ArrayLike, Mat3x3Like, Mat3x3Type from .mat4x4 import Mat4x4, Mat4x4Array, Mat4x4ArrayLike, Mat4x4Like, Mat4x4Type from .point2d import Point2D, Point2DArray, Point2DArrayLike, Point2DLike, Point2DType +from .point3d import Point3D, Point3DArray, Point3DArrayLike, Point3DLike, Point3DType from .quaternion import Quaternion, QuaternionArray, QuaternionArrayLike, QuaternionLike, QuaternionType from .rotation3d import Rotation3D, Rotation3DArray, Rotation3DArrayLike, Rotation3DLike, Rotation3DType from .rotation_axis_angle import ( @@ -118,6 +119,11 @@ "Point2DArrayLike", "Point2DLike", "Point2DType", + "Point3D", + "Point3DArray", + "Point3DArrayLike", + "Point3DLike", + "Point3DType", "Quaternion", "QuaternionArray", "QuaternionArrayLike", diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/point3d.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/point3d.py new file mode 100644 index 000000000000..5fc17487be80 --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/point3d.py @@ -0,0 +1,74 @@ +# NOTE: This file was autogenerated by re_types_builder; DO NOT EDIT. + +from __future__ import annotations + +from typing import Any, Sequence, Tuple, Union + +import numpy.typing as npt +import pyarrow as pa +from attrs import define, field + +from .._baseclasses import ( + BaseExtensionArray, + BaseExtensionType, +) +from ._overrides import point3d_as_array, point3d_native_to_pa_array # noqa: F401 + +__all__ = ["Point3D", "Point3DArray", "Point3DArrayLike", "Point3DLike", "Point3DType"] + + +@define +class Point3D: + """A point in 3D space.""" + + x: float = field(converter=float) + y: float = field(converter=float) + z: float = field(converter=float) + + def __array__(self, dtype: npt.DTypeLike = None) -> npt.NDArray[Any]: + return point3d_as_array(self, dtype=dtype) + + +Point3DLike = Union[Point3D, Sequence[float]] + +Point3DArrayLike = Union[ + Point3D, + Sequence[Point3DLike], + npt.NDArray[Any], + Sequence[npt.NDArray[Any]], + Sequence[Tuple[float, float]], + Sequence[float], +] + + +# --- Arrow support --- + + +class Point3DType(BaseExtensionType): + def __init__(self) -> None: + pa.ExtensionType.__init__( + self, + pa.struct( + [ + pa.field("x", pa.float32(), False, {}), + pa.field("y", pa.float32(), False, {}), + pa.field("z", pa.float32(), False, {}), + ] + ), + "rerun.datatypes.Point3D", + ) + + +class Point3DArray(BaseExtensionArray[Point3DArrayLike]): + _EXTENSION_NAME = "rerun.datatypes.Point3D" + _EXTENSION_TYPE = Point3DType + + @staticmethod + def _native_to_pa_array(data: Point3DArrayLike, data_type: pa.DataType) -> pa.Array: + return point3d_native_to_pa_array(data, data_type) + + +Point3DType._ARRAY_TYPE = Point3DArray + +# TODO(cmc): bring back registration to pyarrow once legacy types are gone +# pa.register_extension_type(Point3DType())