re_components/
mat.rs

1use arrow2::{array::PrimitiveArray, datatypes::DataType};
2use arrow2_convert::{
3    arrow_enable_vec_for_type,
4    deserialize::ArrowDeserialize,
5    field::{ArrowField, FixedSizeVec},
6    serialize::ArrowSerialize,
7};
8
9use super::Vec3D;
10
11/// A 3x3 column-major Matrix made up of 3 Vecs
12///
13/// ```
14/// use re_components::Mat3x3;
15/// use arrow2_convert::field::ArrowField;
16/// use arrow2::datatypes::{DataType, Field};
17///
18/// assert_eq!(
19///     Mat3x3::data_type(),
20///     DataType::FixedSizeList(
21///         Box::new(Field::new("item", DataType::Float32, false)),
22///         9
23///     )
24/// );
25/// ```
26#[derive(Copy, Clone, Debug, PartialEq)]
27#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
28pub struct Mat3x3([Vec3D; 3]);
29
30impl Mat3x3 {
31    pub const IDENTITY: Mat3x3 = Mat3x3([
32        Vec3D([1.0, 0.0, 0.0]),
33        Vec3D([0.0, 1.0, 0.0]),
34        Vec3D([0.0, 0.0, 1.0]),
35    ]);
36}
37
38impl<Idx> std::ops::Index<Idx> for Mat3x3
39where
40    Idx: std::slice::SliceIndex<[Vec3D]>,
41{
42    type Output = Idx::Output;
43
44    #[inline]
45    fn index(&self, index: Idx) -> &Self::Output {
46        &self.0[index]
47    }
48}
49
50impl From<[[f32; 3]; 3]> for Mat3x3 {
51    #[inline]
52    fn from(v: [[f32; 3]; 3]) -> Self {
53        Self([Vec3D(v[0]), Vec3D(v[1]), Vec3D(v[2])])
54    }
55}
56
57#[cfg(feature = "glam")]
58impl From<Mat3x3> for glam::Mat3 {
59    #[inline]
60    fn from(v: Mat3x3) -> Self {
61        Self::from_cols(v[0].into(), v[1].into(), v[2].into())
62    }
63}
64
65#[cfg(feature = "glam")]
66impl From<glam::Mat3> for Mat3x3 {
67    #[inline]
68    fn from(v: glam::Mat3) -> Self {
69        Self::from(v.to_cols_array_2d())
70    }
71}
72
73arrow_enable_vec_for_type!(Mat3x3);
74
75impl ArrowField for Mat3x3 {
76    type Type = Self;
77
78    #[inline]
79    fn data_type() -> DataType {
80        <FixedSizeVec<f32, 9> as ArrowField>::data_type()
81    }
82}
83
84impl ArrowSerialize for Mat3x3 {
85    type MutableArrayType = <FixedSizeVec<f32, 9> as ArrowSerialize>::MutableArrayType;
86
87    #[inline]
88    fn new_array() -> Self::MutableArrayType {
89        FixedSizeVec::<f32, 9>::new_array()
90    }
91
92    #[inline]
93    fn arrow_serialize(v: &Self, array: &mut Self::MutableArrayType) -> arrow2::error::Result<()> {
94        for col in v.0 {
95            array.mut_values().extend_from_slice(&col.0);
96        }
97        array.try_push_valid()
98    }
99}
100
101impl ArrowDeserialize for Mat3x3 {
102    type ArrayType = <FixedSizeVec<f32, 9> as ArrowDeserialize>::ArrayType;
103
104    #[inline]
105    fn arrow_deserialize(
106        v: <&Self::ArrayType as IntoIterator>::Item,
107    ) -> Option<<Self as ArrowField>::Type> {
108        v.map(|v| {
109            let slice = v
110                .as_any()
111                .downcast_ref::<PrimitiveArray<f32>>()
112                .unwrap()
113                .values()
114                .as_slice();
115            Mat3x3([
116                Vec3D(slice[0..3].try_into().unwrap()),
117                Vec3D(slice[3..6].try_into().unwrap()),
118                Vec3D(slice[6..9].try_into().unwrap()),
119            ])
120        })
121    }
122}
123
124#[test]
125fn test_mat3x3_roundtrip() {
126    use arrow2::array::Array;
127    use arrow2_convert::{deserialize::TryIntoCollection, serialize::TryIntoArrow};
128
129    let mats_in: Vec<Mat3x3> = vec![
130        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]].into(),
131        [[11.0, 12.0, 13.0], [14.0, 15.0, 16.0], [17.0, 18.0, 19.0]].into(),
132    ];
133    let array: Box<dyn Array> = mats_in.try_into_arrow().unwrap();
134    let mats_out: Vec<Mat3x3> = TryIntoCollection::try_into_collection(array).unwrap();
135    assert_eq!(mats_in, mats_out);
136}
re_components/mat.rs

re_components/
mat.rs
