1use crate::{SerializationResult, SerializedComponentBatch};
2
3pub trait AsComponents {
20    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch>;
35
36    #[inline]
45    fn to_arrow(
46        &self,
47    ) -> SerializationResult<Vec<(::arrow::datatypes::Field, ::arrow::array::ArrayRef)>> {
48        self.as_serialized_batches()
49            .into_iter()
50            .map(|comp_batch| Ok((arrow::datatypes::Field::from(&comp_batch), comp_batch.array)))
51            .collect()
52    }
53}
54
55#[allow(dead_code)]
56fn assert_object_safe() {
57    let _: &dyn AsComponents;
58}
59
60impl AsComponents for SerializedComponentBatch {
61    #[inline]
62    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch> {
63        vec![self.clone()]
64    }
65}
66
67impl<AS: AsComponents, const N: usize> AsComponents for [AS; N] {
68    #[inline]
69    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch> {
70        self.iter()
71            .flat_map(|as_components| as_components.as_serialized_batches())
72            .collect()
73    }
74}
75
76impl<const N: usize> AsComponents for [&dyn AsComponents; N] {
77    #[inline]
78    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch> {
79        self.iter()
80            .flat_map(|as_components| as_components.as_serialized_batches())
81            .collect()
82    }
83}
84
85impl<const N: usize> AsComponents for [Box<dyn AsComponents>; N] {
86    #[inline]
87    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch> {
88        self.iter()
89            .flat_map(|as_components| as_components.as_serialized_batches())
90            .collect()
91    }
92}
93
94impl<AS: AsComponents> AsComponents for Vec<AS> {
95    #[inline]
96    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch> {
97        self.iter()
98            .flat_map(|as_components| as_components.as_serialized_batches())
99            .collect()
100    }
101}
102
103impl AsComponents for Vec<&dyn AsComponents> {
104    #[inline]
105    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch> {
106        self.iter()
107            .flat_map(|as_components| as_components.as_serialized_batches())
108            .collect()
109    }
110}
111
112impl AsComponents for Vec<Box<dyn AsComponents>> {
113    #[inline]
114    fn as_serialized_batches(&self) -> Vec<SerializedComponentBatch> {
115        self.iter()
116            .flat_map(|as_components| as_components.as_serialized_batches())
117            .collect()
118    }
119}
120
121#[allow(dead_code)]
130#[allow(rustdoc::private_doc_tests)] fn single_ascomponents() {}
132
133#[allow(dead_code)]
138#[allow(rustdoc::private_doc_tests)] fn single_ascomponents_wrapped() {
140    }
143
144#[allow(dead_code)]
149#[allow(rustdoc::private_doc_tests)] fn single_ascomponents_wrapped_many() {
151    }
154
155#[allow(dead_code)]
161#[allow(rustdoc::private_doc_tests)] fn many_ascomponents() {}
163
164#[allow(dead_code)]
170#[allow(rustdoc::private_doc_tests)] fn many_ascomponents_wrapped() {}
172
173#[allow(dead_code)]
179#[allow(rustdoc::private_doc_tests)] fn many_componentbatch_wrapped() {}
181
182#[allow(dead_code)]
188#[allow(rustdoc::private_doc_tests)] fn many_ascomponents_wrapped_many() {}
190
191#[allow(dead_code)]
197#[allow(rustdoc::private_doc_tests)] fn many_componentbatch_wrapped_many() {}
199
200#[cfg(test)]
201mod tests {
202    use std::sync::Arc;
203
204    use arrow::array::{
205        types::UInt32Type, Array as ArrowArray, PrimitiveArray as ArrowPrimitiveArray,
206    };
207    use itertools::Itertools as _;
208    use similar_asserts::assert_eq;
209
210    use crate::{Component as _, ComponentDescriptor};
211
212    #[derive(Clone, Copy, Debug, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)]
213    #[repr(transparent)]
214    pub struct MyColor(pub u32);
215
216    impl MyColor {
217        fn descriptor() -> ComponentDescriptor {
218            ComponentDescriptor {
219                archetype: Some("test".into()),
220                component: "color".into(),
221                component_type: Some(Self::name()),
222            }
223        }
224    }
225
226    crate::macros::impl_into_cow!(MyColor);
227
228    impl re_byte_size::SizeBytes for MyColor {
229        #[inline]
230        fn heap_size_bytes(&self) -> u64 {
231            let Self(_) = self;
232            0
233        }
234    }
235
236    impl crate::Loggable for MyColor {
237        fn arrow_datatype() -> arrow::datatypes::DataType {
238            arrow::datatypes::DataType::UInt32
239        }
240
241        fn to_arrow_opt<'a>(
242            data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
243        ) -> crate::SerializationResult<arrow::array::ArrayRef>
244        where
245            Self: 'a,
246        {
247            use crate::datatypes::UInt32;
248            UInt32::to_arrow_opt(
249                data.into_iter()
250                    .map(|opt| opt.map(Into::into).map(|c| UInt32(c.0))),
251            )
252        }
253
254        fn from_arrow_opt(
255            data: &dyn arrow::array::Array,
256        ) -> crate::DeserializationResult<Vec<Option<Self>>> {
257            use crate::datatypes::UInt32;
258            Ok(UInt32::from_arrow_opt(data)?
259                .into_iter()
260                .map(|opt| opt.map(|v| Self(v.0)))
261                .collect())
262        }
263    }
264
265    impl crate::Component for MyColor {
266        fn name() -> crate::ComponentType {
267            "example.MyColor".into()
268        }
269    }
270
271    #[allow(dead_code)]
272    fn data() -> (MyColor, MyColor, MyColor, Vec<MyColor>) {
273        let red = MyColor(0xDD0000FF);
274        let green = MyColor(0x00DD00FF);
275        let blue = MyColor(0x0000DDFF);
276        let colors = vec![red, green, blue];
277        (red, green, blue, colors)
278    }
279
280    #[test]
281    fn single_ascomponents_howto() {
282        let (red, _, _, _) = data();
283
284        let got = {
285            let red = &red as &dyn crate::ComponentBatch;
286            vec![red.try_serialized(MyColor::descriptor()).unwrap().array]
287        };
288        let expected = vec![
289            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![red.0])) as Arc<dyn ArrowArray>,
290        ];
291        assert_eq!(&expected, &got);
292    }
293
294    #[test]
295    fn single_componentbatch() -> anyhow::Result<()> {
296        let (red, _, _, _) = data();
297
298        let got = (&red as &dyn crate::ComponentBatch).to_arrow()?;
300        let expected =
301            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![red.0])) as Arc<dyn ArrowArray>;
302        similar_asserts::assert_eq!(&expected, &got);
303
304        Ok(())
305    }
306
307    #[test]
308    fn single_ascomponents_wrapped_howto() {
309        let (red, _, _, _) = data();
310
311        let got = {
312            let red = &red as &dyn crate::ComponentBatch;
313            vec![red.try_serialized(MyColor::descriptor()).unwrap().array]
314        };
315        let expected = vec![
316            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![red.0])) as Arc<dyn ArrowArray>,
317        ];
318        assert_eq!(&expected, &got);
319    }
320
321    #[test]
322    fn single_componentbatch_wrapped() -> anyhow::Result<()> {
323        let (red, _, _, _) = data();
324
325        let got = (&[red] as &dyn crate::ComponentBatch).to_arrow()?;
327        let expected =
328            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![red.0])) as Arc<dyn ArrowArray>;
329        similar_asserts::assert_eq!(&expected, &got);
330
331        Ok(())
332    }
333
334    #[test]
335    fn single_ascomponents_wrapped_many_howto() {
336        let (red, green, blue, _) = data();
337
338        let got = {
339            let red = &red as &dyn crate::ComponentBatch;
340            let green = &green as &dyn crate::ComponentBatch;
341            let blue = &blue as &dyn crate::ComponentBatch;
342            [
343                red.try_serialized(MyColor::descriptor()).unwrap(),
344                green.try_serialized(MyColor::descriptor()).unwrap(),
345                blue.try_serialized(MyColor::descriptor()).unwrap(),
346            ]
347            .into_iter()
348            .map(|batch| batch.array)
349            .collect_vec()
350        };
351        let expected = vec![
352            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![red.0])) as Arc<dyn ArrowArray>,
353            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![green.0])) as Arc<dyn ArrowArray>,
354            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![blue.0])) as Arc<dyn ArrowArray>,
355        ];
356        assert_eq!(&expected, &got);
357    }
358
359    #[test]
360    fn single_componentbatch_wrapped_many() -> anyhow::Result<()> {
361        let (red, green, blue, _) = data();
362
363        let got = (&[red, green, blue] as &dyn crate::ComponentBatch).to_arrow()?;
365        let expected = Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![
366            red.0, green.0, blue.0,
367        ])) as Arc<dyn ArrowArray>;
368        similar_asserts::assert_eq!(&expected, &got);
369
370        Ok(())
371    }
372
373    #[test]
374    fn many_componentbatch() -> anyhow::Result<()> {
375        let (red, green, blue, colors) = data();
376
377        let got = (&colors as &dyn crate::ComponentBatch).to_arrow()?;
379        let expected = Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![
380            red.0, green.0, blue.0,
381        ])) as Arc<dyn ArrowArray>;
382        similar_asserts::assert_eq!(&expected, &got);
383
384        Ok(())
385    }
386
387    #[test]
388    fn many_ascomponents_wrapped_howto() {
389        let (red, green, blue, colors) = data();
390
391        let got = {
392            let colors = &colors as &dyn crate::ComponentBatch;
393            vec![colors.try_serialized(MyColor::descriptor()).unwrap().array]
394        };
395        let expected = vec![Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![
396            red.0, green.0, blue.0,
397        ])) as Arc<dyn ArrowArray>];
398        assert_eq!(&expected, &got);
399    }
400
401    #[test]
402    fn many_ascomponents_wrapped_many_howto() {
403        let (red, green, blue, colors) = data();
404
405        let got = {
407            let colors = &colors as &dyn crate::ComponentBatch;
408            vec![
409                colors.try_serialized(MyColor::descriptor()).unwrap().array,
410                colors.try_serialized(MyColor::descriptor()).unwrap().array,
411                colors.try_serialized(MyColor::descriptor()).unwrap().array,
412            ]
413        };
414        let expected = vec![
415            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![
416                red.0, green.0, blue.0,
417            ])) as Arc<dyn ArrowArray>,
418            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![
419                red.0, green.0, blue.0,
420            ])) as Arc<dyn ArrowArray>,
421            Arc::new(ArrowPrimitiveArray::<UInt32Type>::from(vec![
422                red.0, green.0, blue.0,
423            ])) as Arc<dyn ArrowArray>,
424        ];
425        assert_eq!(&expected, &got);
426    }
427}