narrow/array/
struct.rs

1//! Array for product types.
2
3use super::{Array, ArrayType};
4use crate::{
5    Length,
6    bitmap::{Bitmap, BitmapRef, BitmapRefMut, ValidityBitmap},
7    buffer::{BufferType, VecBuffer},
8    nullability::{NonNullable, Nullability, Nullable},
9    validity::Validity,
10};
11
12/// Struct array types.
13pub trait StructArrayType: ArrayType<Self> {
14    /// The array type that stores items of this struct. Note this differs from
15    /// the [`ArrayType`] array because that wraps this array. Also note that this
16    /// has no [`Array`] bound.
17    // TODO(mbrobbe): add offset and union generics
18    type Array<Buffer: BufferType>; // into<fields> this then requires all arraytype impls to provide a field
19}
20
21/// Array for product types.
22pub struct StructArray<
23    T: StructArrayType,
24    Nullable: Nullability = NonNullable,
25    Buffer: BufferType = VecBuffer,
26>(pub Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer>);
27
28impl<T: StructArrayType, Nullable: Nullability, Buffer: BufferType> StructArray<T, Nullable, Buffer>
29where
30    for<'a> &'a Self: IntoIterator,
31{
32    /// Returns an iterator over the items in this [`StructArray`].
33    pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
34        self.into_iter()
35    }
36}
37
38impl<T: StructArrayType, Nullable: Nullability, Buffer: BufferType> Array
39    for StructArray<T, Nullable, Buffer>
40{
41    type Item = Nullable::Item<T>;
42}
43
44impl<T: StructArrayType, Nullable: Nullability, Buffer: BufferType> Clone
45    for StructArray<T, Nullable, Buffer>
46where
47    Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer>: Clone,
48{
49    fn clone(&self) -> Self {
50        Self(self.0.clone())
51    }
52}
53
54impl<T: StructArrayType, Nullable: Nullability, Buffer: BufferType> Default
55    for StructArray<T, Nullable, Buffer>
56where
57    Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer>: Default,
58{
59    fn default() -> Self {
60        Self(Default::default())
61    }
62}
63
64impl<T: StructArrayType, Buffer: BufferType> From<StructArray<T, NonNullable, Buffer>>
65    for StructArray<T, Nullable, Buffer>
66where
67    <T as StructArrayType>::Array<Buffer>: Length,
68    Bitmap<Buffer>: FromIterator<bool>,
69{
70    fn from(value: StructArray<T, NonNullable, Buffer>) -> Self {
71        Self(Validity::from(value.0))
72    }
73}
74
75impl<T: StructArrayType, U, Nullable: Nullability, Buffer: BufferType> FromIterator<U>
76    for StructArray<T, Nullable, Buffer>
77where
78    Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer>: FromIterator<U>,
79{
80    fn from_iter<I: IntoIterator<Item = U>>(iter: I) -> Self {
81        Self(iter.into_iter().collect())
82    }
83}
84
85impl<T: StructArrayType, U, Buffer: BufferType> Extend<U> for StructArray<T, NonNullable, Buffer>
86where
87    <T as StructArrayType>::Array<Buffer>: Extend<U>,
88{
89    fn extend<I: IntoIterator<Item = U>>(&mut self, iter: I) {
90        self.0.extend(iter);
91    }
92}
93
94impl<T: StructArrayType, U, Buffer: BufferType> Extend<Option<U>>
95    for StructArray<T, Nullable, Buffer>
96where
97    Validity<<T as StructArrayType>::Array<Buffer>, Buffer>: Extend<Option<U>>,
98{
99    fn extend<I: IntoIterator<Item = Option<U>>>(&mut self, iter: I) {
100        self.0.extend(iter);
101    }
102}
103
104impl<T: StructArrayType, Nullable: Nullability, Buffer: BufferType> IntoIterator
105    for StructArray<T, Nullable, Buffer>
106where
107    Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer>: IntoIterator,
108{
109    type Item =
110        <Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer> as IntoIterator>::Item;
111    type IntoIter = <Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer> as
112        IntoIterator>::IntoIter;
113
114    fn into_iter(self) -> Self::IntoIter {
115        self.0.into_iter()
116    }
117}
118
119impl<'a, T: StructArrayType, Nullable: Nullability, Buffer: BufferType> IntoIterator
120    for &'a StructArray<T, Nullable, Buffer>
121where
122    &'a Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer>: IntoIterator,
123{
124    type Item = <&'a Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer> as
125        IntoIterator>::Item;
126    type IntoIter = <&'a Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer> as IntoIterator>::IntoIter;
127
128    fn into_iter(self) -> Self::IntoIter {
129        self.0.into_iter()
130    }
131}
132
133impl<T: StructArrayType, Nullable: Nullability, Buffer: BufferType> Length
134    for StructArray<T, Nullable, Buffer>
135where
136    Nullable::Collection<<T as StructArrayType>::Array<Buffer>, Buffer>: Length,
137{
138    fn len(&self) -> usize {
139        self.0.len()
140    }
141}
142
143impl<T: StructArrayType, Buffer: BufferType> BitmapRef for StructArray<T, Nullable, Buffer> {
144    type Buffer = Buffer;
145
146    fn bitmap_ref(&self) -> &Bitmap<Self::Buffer> {
147        self.0.bitmap_ref()
148    }
149}
150
151impl<T: StructArrayType, Buffer: BufferType> BitmapRefMut for StructArray<T, Nullable, Buffer> {
152    fn bitmap_ref_mut(&mut self) -> &mut Bitmap<Self::Buffer> {
153        self.0.bitmap_ref_mut()
154    }
155}
156
157impl<T: StructArrayType, Buffer: BufferType> ValidityBitmap for StructArray<T, Nullable, Buffer> {}
158
159#[cfg(test)]
160mod tests {
161    use crate::{
162        array::{ArrayTypeOf, OptionArrayTypeOf, UnionType},
163        offset::Offset,
164    };
165
166    use super::*;
167
168    // Definition
169    #[derive(Default)]
170    struct Foo<'a> {
171        a: u32,
172        b: Option<()>,
173        c: (),
174        d: Option<[u64; 2]>,
175        e: bool,
176        f: &'a [u8],
177        g: String,
178    }
179    // These impls below can all be generated.
180    impl<'a> ArrayType<Self> for Foo<'a> {
181        type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
182            StructArray<Foo<'a>, NonNullable, Buffer>;
183    }
184    impl<'a> ArrayType<Foo<'a>> for Option<Foo<'a>> {
185        type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
186            StructArray<Foo<'a>, Nullable, Buffer>;
187    }
188
189    struct FooArray<'a, Buffer: BufferType> {
190        a: ArrayTypeOf<u32, Buffer>,
191        b: OptionArrayTypeOf<(), Buffer>,
192        c: ArrayTypeOf<(), Buffer>,
193        d: OptionArrayTypeOf<[u64; 2], Buffer>,
194        e: ArrayTypeOf<bool, Buffer>,
195        f: ArrayTypeOf<&'a [u8], Buffer>,
196        g: ArrayTypeOf<String, Buffer>,
197    }
198
199    impl<'a, Buffer: BufferType> Default for FooArray<'a, Buffer>
200    where
201        ArrayTypeOf<u32, Buffer>: Default,
202        OptionArrayTypeOf<(), Buffer>: Default,
203        ArrayTypeOf<(), Buffer>: Default,
204        OptionArrayTypeOf<[u64; 2], Buffer>: Default,
205        ArrayTypeOf<bool, Buffer>: Default,
206        ArrayTypeOf<&'a [u8], Buffer>: Default,
207        ArrayTypeOf<String, Buffer>: Default,
208    {
209        fn default() -> Self {
210            Self {
211                a: <ArrayTypeOf<u32, Buffer>>::default(),
212                b: <OptionArrayTypeOf<(), Buffer>>::default(),
213                c: <ArrayTypeOf<(), Buffer>>::default(),
214                d: <OptionArrayTypeOf<[u64; 2], Buffer>>::default(),
215                e: <ArrayTypeOf<bool, Buffer>>::default(),
216                f: <ArrayTypeOf<&'a [u8], Buffer>>::default(),
217                g: <ArrayTypeOf<String, Buffer>>::default(),
218            }
219        }
220    }
221
222    impl<'a, Buffer: BufferType> Extend<Foo<'a>> for FooArray<'a, Buffer>
223    where
224        ArrayTypeOf<u32, Buffer>: Extend<u32>,
225        OptionArrayTypeOf<(), Buffer>: Extend<Option<()>>,
226        ArrayTypeOf<(), Buffer>: Extend<()>,
227        OptionArrayTypeOf<[u64; 2], Buffer>: Extend<Option<[u64; 2]>>,
228        ArrayTypeOf<bool, Buffer>: Extend<bool>,
229        ArrayTypeOf<&'a [u8], Buffer>: Extend<&'a [u8]>,
230        ArrayTypeOf<String, Buffer>: Extend<String>,
231    {
232        fn extend<I: IntoIterator<Item = Foo<'a>>>(&mut self, iter: I) {
233            iter.into_iter().for_each(
234                |Foo {
235                     a,
236                     b,
237                     c,
238                     d,
239                     e,
240                     f,
241                     g,
242                 }| {
243                    self.a.extend(std::iter::once(a));
244                    self.b.extend(std::iter::once(b));
245                    self.c.extend(std::iter::once(c));
246                    self.d.extend(std::iter::once(d));
247                    self.e.extend(std::iter::once(e));
248                    self.f.extend(std::iter::once(f));
249                    self.g.extend(std::iter::once(g));
250                },
251            );
252        }
253    }
254
255    impl<'a, Buffer: BufferType> FromIterator<Foo<'a>> for FooArray<'a, Buffer>
256    where
257        ArrayTypeOf<u32, Buffer>: Default + Extend<u32>,
258        OptionArrayTypeOf<(), Buffer>: Default + Extend<Option<()>>,
259        ArrayTypeOf<(), Buffer>: Default + Extend<()>,
260        OptionArrayTypeOf<[u64; 2], Buffer>: Default + Extend<Option<[u64; 2]>>,
261        ArrayTypeOf<bool, Buffer>: Default + Extend<bool>,
262        ArrayTypeOf<&'a [u8], Buffer>: Default + Extend<&'a [u8]>,
263        ArrayTypeOf<String, Buffer>: Default + Extend<String>,
264    {
265        #[allow(clippy::many_single_char_names)]
266        fn from_iter<T: IntoIterator<Item = Foo<'a>>>(iter: T) -> Self {
267            let (a, (b, (c, (d, (e, (f, g)))))) = iter
268                .into_iter()
269                .map(
270                    |Foo {
271                         a,
272                         b,
273                         c,
274                         d,
275                         e,
276                         f,
277                         g,
278                     }| (a, (b, (c, (d, (e, (f, g)))))),
279                )
280                .unzip();
281            Self {
282                a,
283                b,
284                c,
285                d,
286                e,
287                f,
288                g,
289            }
290        }
291    }
292    impl<'a> StructArrayType for Foo<'a> {
293        type Array<Buffer: BufferType> = FooArray<'a, Buffer>;
294    }
295
296    impl<Buffer: BufferType> Length for FooArray<'_, Buffer>
297    where
298        ArrayTypeOf<u32, Buffer>: Length,
299    {
300        fn len(&self) -> usize {
301            self.a.len()
302        }
303    }
304
305    #[test]
306    fn from_iter() {
307        let input = [
308            Foo {
309                a: 1,
310                b: None,
311                c: (),
312                d: Some([1, 2]),
313                e: false,
314                f: &[1],
315                g: "a".to_owned(),
316            },
317            Foo {
318                a: 2,
319                b: Some(()),
320                c: (),
321                d: Some([3, 4]),
322                e: true,
323                f: &[2, 3],
324                g: "s".to_owned(),
325            },
326            Foo {
327                a: 3,
328                b: None,
329                c: (),
330                d: None,
331                e: true,
332                f: &[4],
333                g: "d".to_owned(),
334            },
335            Foo {
336                a: 4,
337                b: None,
338                c: (),
339                d: None,
340                e: true,
341                f: &[],
342                g: "f".to_owned(),
343            },
344        ];
345        let array = input.into_iter().collect::<StructArray<Foo>>();
346        assert_eq!(array.len(), 4);
347        assert_eq!(array.0.a.into_iter().collect::<Vec<_>>(), &[1, 2, 3, 4]);
348        assert_eq!(
349            array.0.b.into_iter().collect::<Vec<_>>(),
350            &[None, Some(()), None, None]
351        );
352        assert_eq!(array.0.c.into_iter().collect::<Vec<_>>(), &[(), (), (), ()]);
353        assert_eq!(
354            array.0.d.into_iter().collect::<Vec<_>>(),
355            &[Some([1, 2]), Some([3, 4]), None, None]
356        );
357        assert_eq!(
358            array.0.e.into_iter().collect::<Vec<_>>(),
359            &[false, true, true, true]
360        );
361        assert_eq!(
362            array.0.f.0.data.into_iter().collect::<Vec<_>>(),
363            &[1, 2, 3, 4]
364        );
365        assert_eq!(
366            array.0.f.0.offsets.into_iter().collect::<Vec<_>>(),
367            &[0, 1, 3, 4, 4]
368        );
369        assert_eq!(
370            array.0.g.0.0.data.into_iter().collect::<Vec<_>>(),
371            &[97, 115, 100, 102] // a s d f
372        );
373        assert_eq!(
374            array.0.g.0.0.offsets.into_iter().collect::<Vec<_>>(),
375            &[0, 1, 2, 3, 4]
376        );
377
378        let input_nullable = [
379            None,
380            Some(Foo {
381                a: 1,
382                b: None,
383                c: (),
384                d: Some([1, 2]),
385                e: false,
386                f: &[1],
387                g: "a".to_owned(),
388            }),
389        ];
390        let array_nullable = input_nullable
391            .into_iter()
392            .collect::<StructArray<Foo, Nullable>>();
393        assert_eq!(array_nullable.len(), 2);
394        assert_eq!(array_nullable.is_null(0), Some(true));
395        assert_eq!(array_nullable.is_valid(1), Some(true));
396        assert_eq!(array_nullable.is_valid(2), None);
397    }
398
399    #[cfg(feature = "derive")]
400    #[test]
401    fn derive() {
402        #[derive(crate::ArrayType, Copy, Clone, Default)]
403        struct Unit;
404
405        #[derive(crate::ArrayType)]
406        struct Foo(Option<Bar<u32>>);
407
408        #[derive(crate::ArrayType)]
409        struct Bar<T>(T);
410
411        #[derive(crate::ArrayType, Default)]
412        struct FooBar(Option<Vec<Option<u32>>>);
413
414        let mut foo_bar = StructArray::<FooBar, NonNullable>::default();
415        foo_bar.extend(std::iter::once(FooBar(None)));
416        foo_bar.extend(std::iter::once(FooBar(Some(vec![None]))));
417        let mut foo_bar_nullable = StructArray::<FooBar, Nullable>::default();
418        foo_bar_nullable.extend(std::iter::once(Some(FooBar(None))));
419        foo_bar_nullable.extend(std::iter::once(None));
420    }
421
422    #[cfg(feature = "derive")]
423    #[test]
424    fn into_iter() {
425        #[derive(crate::ArrayType, Copy, Clone, Debug, Default, PartialEq)]
426        struct Unit;
427
428        #[derive(crate::ArrayType, Copy, Clone, Debug, Default, PartialEq)]
429        struct Unnamed(u8, Option<u16>, u32, u64);
430
431        #[derive(crate::ArrayType, Clone, Debug, Default, PartialEq)]
432        struct Named {
433            a: u8,
434            b: bool,
435            c: Option<u16>,
436            d: Option<bool>,
437            e: String,
438            f: Option<String>,
439        }
440
441        let unit_input = [Unit; 3];
442        let unit_array = unit_input.into_iter().collect::<StructArray<Unit>>();
443        assert_eq!(unit_array.len(), unit_input.len());
444        let unit_output = unit_array.into_iter().collect::<Vec<_>>();
445        assert_eq!(unit_output, unit_input);
446        let unit_input_nullable = unit_input.map(Option::Some);
447        let unit_array_nullable = unit_input_nullable
448            .into_iter()
449            .collect::<StructArray<Unit, Nullable>>();
450        assert_eq!(unit_array_nullable.len(), unit_input_nullable.len());
451        let unit_output_nullable = unit_array_nullable.into_iter().collect::<Vec<_>>();
452        assert_eq!(unit_output_nullable, unit_input_nullable);
453
454        let unnamed_input = [Unnamed(1, Some(2), 3, 4); 3];
455        let unnamed_array = unnamed_input.into_iter().collect::<StructArray<Unnamed>>();
456        assert_eq!(unnamed_array.len(), unnamed_input.len());
457        let unnamed_output = unnamed_array.into_iter().collect::<Vec<_>>();
458        assert_eq!(unnamed_output, unnamed_input);
459        let unnamed_input_nullable = unnamed_input.map(Option::Some);
460        let unnamed_array_nullable = unnamed_input_nullable
461            .into_iter()
462            .collect::<StructArray<Unnamed, Nullable>>();
463        assert_eq!(unnamed_array_nullable.len(), unnamed_input_nullable.len());
464        let unnamed_output_nullable = unnamed_array_nullable.into_iter().collect::<Vec<_>>();
465        assert_eq!(unnamed_output_nullable, unnamed_input_nullable);
466
467        let named_input = [Named {
468            a: 1,
469            b: false,
470            c: Some(3),
471            d: Some(true),
472            e: "hello".to_owned(),
473            f: None,
474        }];
475        let named_array = named_input
476            .clone()
477            .into_iter()
478            .collect::<StructArray<Named>>();
479        assert_eq!(named_array.len(), named_input.len());
480        let named_output = named_array.into_iter().collect::<Vec<_>>();
481        assert_eq!(named_output, named_input);
482        let named_input_nullable = named_input
483            .into_iter()
484            .map(Option::Some)
485            .collect::<Vec<_>>();
486        let named_array_nullable = named_input_nullable
487            .clone()
488            .into_iter()
489            .collect::<StructArray<Named, Nullable>>();
490        assert_eq!(named_array_nullable.len(), named_input_nullable.len());
491        let named_output_nullable = named_array_nullable.into_iter().collect::<Vec<_>>();
492        assert_eq!(named_output_nullable, named_input_nullable);
493    }
494
495    #[cfg(feature = "derive")]
496    #[test]
497    fn nested_option_derived() {
498        use std::collections::VecDeque;
499
500        {
501            #[derive(crate::ArrayType, Clone, Debug, PartialEq)]
502            struct Foo(Vec<Option<String>>);
503
504            let input = [Foo(vec![None]), Foo(vec![Some("hello".to_owned())])];
505            let array = input.clone().into_iter().collect::<StructArray<Foo>>();
506            assert_eq!(array.len(), 2);
507            let output = array.into_iter().collect::<Vec<_>>();
508            assert_eq!(input.as_slice(), output);
509        };
510
511        {
512            #[derive(crate::ArrayType, Clone, Debug, PartialEq)]
513            struct Foo([Option<String>; 1]);
514
515            let input = [Foo([None]), Foo([Some("hello".to_owned())])];
516            let array = input.clone().into_iter().collect::<StructArray<Foo>>();
517            assert_eq!(array.len(), 2);
518            let output = array.into_iter().collect::<Vec<_>>();
519            assert_eq!(input.as_slice(), output);
520        };
521
522        {
523            #[derive(crate::ArrayType, Clone, Debug, PartialEq)]
524            struct Foo(VecDeque<Option<String>>);
525
526            let input = [
527                Foo(VecDeque::from_iter([None])),
528                Foo(VecDeque::from_iter([Some("hello".to_owned())])),
529            ];
530            let array = input.clone().into_iter().collect::<StructArray<Foo>>();
531            assert_eq!(array.len(), 2);
532            // TODO(mbrobbel): add support to offset iterator for VecDeque
533            // OR fix bound + convert via proc macro
534            // let output = array.into_iter().collect::<Vec<_>>();
535            // assert_eq!(input.as_slice(), output);
536        };
537    }
538}