narrow/array/
union.rs

1//! Array for sum types.
2
3use std::iter;
4
5use crate::{
6    Length, NonNullable,
7    buffer::{BufferType, VecBuffer},
8    offset::{self, Offset},
9};
10
11use super::{Array, ArrayType, Int8Array, Int32Array};
12
13/// Different types of union layouts.
14pub trait UnionType {
15    /// The array for this union type.
16    type Array<
17        T: UnionArrayType<VARIANTS>,
18        const VARIANTS: usize,
19        Buffer: BufferType,
20        OffsetItem: Offset,
21    >
22    where
23        for<'a> i8: From<&'a T>;
24}
25
26/// The dense union layout.
27#[derive(Clone, Copy, Debug)]
28pub struct DenseLayout;
29
30impl UnionType for DenseLayout {
31    type Array<
32        T: UnionArrayType<VARIANTS>,
33        const VARIANTS: usize,
34        Buffer: BufferType,
35        OffsetItem: Offset,
36    >
37        = DenseUnionArray<T, VARIANTS, Buffer, OffsetItem>
38    where
39        for<'a> i8: From<&'a T>;
40}
41
42/// The sparse union layout.
43#[derive(Clone, Copy, Debug)]
44pub struct SparseLayout;
45
46impl UnionType for SparseLayout {
47    type Array<
48        T: UnionArrayType<VARIANTS>,
49        const VARIANTS: usize,
50        Buffer: BufferType,
51        OffsetItem: Offset,
52    >
53        = SparseUnionArray<T, VARIANTS, Buffer, OffsetItem>
54    where
55        for<'a> i8: From<&'a T>;
56}
57
58/// Indicates that a [`UnionType`] generic is not applicable.
59///
60/// This is used instead to prevent confusion in code because we don't have default
61/// types for generic associated types.
62///
63/// This still shows up as [`SparseLayout`] in documentation but there is no way
64/// to prevent that.
65pub type NA = SparseLayout;
66
67/// Union array types.
68pub trait UnionArrayType<const VARIANTS: usize>
69where
70    for<'a> i8: From<&'a Self>,
71{
72    // can't use this yet in const expressions, unfortunately
73    /// The number of variants.
74    const VARIANTS: usize = VARIANTS;
75
76    /// The array type storing the variants of the union array.
77    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType>;
78}
79
80/// The array data for enum variants stored in union array wrappers.
81///
82/// Implementations provide the method to convert back to the original enum.
83pub trait EnumVariant<const INDEX: usize>: Sized {
84    /// The data for this variant. It must be an `ArrayType` because it is stored in an array.
85    /// And it must implement `Into<Self>` (this is taking the data and wrapping it in the original enum variant).
86    type Data: ArrayType<Self::Data> + Default;
87
88    /// Wraps the data in the original enum variant
89    fn from_data(value: Self::Data) -> Self;
90}
91
92/// Array for sum types.
93pub struct UnionArray<
94    T: UnionArrayType<VARIANTS>,
95    // we need this const here because:
96    // generic parameters may not be used in const operations
97    // type parameters may not be used in const expressions
98    const VARIANTS: usize,
99    UnionLayout: UnionType = DenseLayout,
100    Buffer: BufferType = VecBuffer,
101    OffsetItem: Offset = offset::NA,
102>(pub(crate) UnionLayout::Array<T, VARIANTS, Buffer, OffsetItem>)
103where
104    for<'a> i8: From<&'a T>;
105
106impl<
107    T: UnionArrayType<VARIANTS>,
108    const VARIANTS: usize,
109    UnionLayout: UnionType,
110    Buffer: BufferType,
111    OffsetItem: Offset,
112> Array for UnionArray<T, VARIANTS, UnionLayout, Buffer, OffsetItem>
113where
114    for<'a> i8: From<&'a T>,
115{
116    type Item = T;
117}
118
119impl<
120    T: UnionArrayType<VARIANTS>,
121    const VARIANTS: usize,
122    UnionLayout: UnionType,
123    Buffer: BufferType,
124    OffsetItem: Offset,
125> Clone for UnionArray<T, VARIANTS, UnionLayout, Buffer, OffsetItem>
126where
127    for<'a> i8: From<&'a T>,
128    UnionLayout::Array<T, VARIANTS, Buffer, OffsetItem>: Clone,
129{
130    fn clone(&self) -> Self {
131        Self(self.0.clone())
132    }
133}
134
135impl<
136    T: UnionArrayType<VARIANTS>,
137    const VARIANTS: usize,
138    UnionLayout: UnionType,
139    Buffer: BufferType,
140    OffsetItem: Offset,
141> Default for UnionArray<T, VARIANTS, UnionLayout, Buffer, OffsetItem>
142where
143    for<'a> i8: From<&'a T>,
144    UnionLayout::Array<T, VARIANTS, Buffer, OffsetItem>: Default,
145{
146    fn default() -> Self {
147        Self(Default::default())
148    }
149}
150
151impl<
152    T: UnionArrayType<VARIANTS>,
153    const VARIANTS: usize,
154    UnionLayout: UnionType,
155    Buffer: BufferType,
156    OffsetItem: Offset,
157> Extend<T> for UnionArray<T, VARIANTS, UnionLayout, Buffer, OffsetItem>
158where
159    for<'a> i8: From<&'a T>,
160    <UnionLayout as UnionType>::Array<T, VARIANTS, Buffer, OffsetItem>: Extend<T>,
161{
162    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
163        self.0.extend(iter);
164    }
165}
166
167impl<
168    T: UnionArrayType<VARIANTS>,
169    const VARIANTS: usize,
170    UnionLayout: UnionType,
171    Buffer: BufferType,
172    OffsetItem: Offset,
173> Length for UnionArray<T, VARIANTS, UnionLayout, Buffer, OffsetItem>
174where
175    for<'a> i8: From<&'a T>,
176    <UnionLayout as UnionType>::Array<T, VARIANTS, Buffer, OffsetItem>: Length,
177{
178    fn len(&self) -> usize {
179        self.0.len()
180    }
181}
182
183impl<
184    T: UnionArrayType<VARIANTS>,
185    const VARIANTS: usize,
186    UnionLayout: UnionType,
187    Buffer: BufferType,
188    OffsetItem: Offset,
189> FromIterator<T> for UnionArray<T, VARIANTS, UnionLayout, Buffer, OffsetItem>
190where
191    for<'a> i8: From<&'a T>,
192    OffsetItem: Offset,
193    <UnionLayout as UnionType>::Array<T, VARIANTS, Buffer, OffsetItem>: FromIterator<T>,
194{
195    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
196        Self(iter.into_iter().collect())
197    }
198}
199
200/// A dense union array.
201pub struct DenseUnionArray<
202    T: UnionArrayType<VARIANTS>,
203    const VARIANTS: usize,
204    Buffer: BufferType = VecBuffer,
205    OffsetItem: Offset = i32,
206> where
207    for<'a> i8: From<&'a T>,
208{
209    /// The data for the variants
210    pub variants: <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>,
211    /// The types field encodes the variants
212    pub types: Int8Array<NonNullable, Buffer>,
213    /// The offsets in the variant arrays
214    pub offsets: Int32Array<NonNullable, Buffer>,
215}
216
217/// A trait that should be implemented by the derive macro for dense layout
218/// union arrays.
219#[doc(hidden)]
220pub trait DenseOffset {
221    /// Returns the length (number of items) of the variant with the given type
222    /// id stored in this array.
223    fn variant_len(&self, type_id: i8) -> usize;
224}
225
226impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
227    Clone for DenseUnionArray<T, VARIANTS, Buffer, OffsetItem>
228where
229    for<'a> i8: From<&'a T>,
230    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>: Clone,
231    Int8Array<NonNullable, Buffer>: Clone,
232    Int32Array<NonNullable, Buffer>: Clone,
233{
234    fn clone(&self) -> Self {
235        Self {
236            variants: self.variants.clone(),
237            types: self.types.clone(),
238            offsets: self.offsets.clone(),
239        }
240    }
241}
242
243impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
244    Default for DenseUnionArray<T, VARIANTS, Buffer, OffsetItem>
245where
246    for<'a> i8: From<&'a T>,
247    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>: Default,
248    Int8Array<NonNullable, Buffer>: Default,
249    Int32Array<NonNullable, Buffer>: Default,
250{
251    fn default() -> Self {
252        Self {
253            variants: Default::default(),
254            types: Int8Array::default(),
255            offsets: Int32Array::default(),
256        }
257    }
258}
259
260impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
261    Length for DenseUnionArray<T, VARIANTS, Buffer, OffsetItem>
262where
263    for<'a> i8: From<&'a T>,
264{
265    fn len(&self) -> usize {
266        self.types.len()
267    }
268}
269
270impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
271    Extend<T> for DenseUnionArray<T, VARIANTS, Buffer, OffsetItem>
272where
273    for<'a> i8: From<&'a T>,
274    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>:
275        Extend<T> + DenseOffset,
276    Int8Array<NonNullable, Buffer>: Extend<i8>,
277    Int32Array<NonNullable, Buffer>: Extend<i32>,
278{
279    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
280        iter.into_iter().for_each(|item| {
281            let type_id = i8::from(&item);
282            let idx = usize::try_from(type_id).expect("bad type id");
283            assert!(idx < VARIANTS, "type id greater than number of variants");
284
285            // For dense unions, we need to track the current offset for each variant
286            // Count the current elements of this type
287            let offset = self.variants.variant_len(type_id);
288
289            self.types.extend(iter::once(type_id));
290            self.offsets
291                .extend(iter::once(i32::try_from(offset).expect("overflow")));
292            self.variants.extend(iter::once(item));
293        });
294    }
295}
296
297impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
298    FromIterator<T> for DenseUnionArray<T, VARIANTS, Buffer, OffsetItem>
299where
300    for<'a> i8: From<&'a T>,
301    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>: Default + Extend<T>,
302    <Buffer as BufferType>::Buffer<i8>: Default + Extend<i8>,
303    <Buffer as BufferType>::Buffer<i32>: Default + Extend<i32>,
304{
305    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
306        let mut lens = [0; VARIANTS];
307        let ((types, offsets), variants) = iter
308            .into_iter()
309            .map(|item| {
310                let type_id = i8::from(&item);
311                let idx = usize::try_from(type_id).expect("bad type id");
312
313                assert!(idx < VARIANTS, "type id greater than number of variants");
314                let result = ((type_id, lens[idx]), item);
315
316                lens[idx] += 1;
317                result
318            })
319            .unzip();
320        Self {
321            variants,
322            types,
323            offsets,
324        }
325    }
326}
327
328/// A sparse union array.
329pub struct SparseUnionArray<
330    T: UnionArrayType<VARIANTS>,
331    const VARIANTS: usize,
332    Buffer: BufferType = VecBuffer,
333    OffsetItem: Offset = i32,
334> where
335    for<'a> i8: From<&'a T>,
336{
337    /// The data for the variants
338    pub variants: <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>,
339    /// The types field encodes the variants
340    pub types: Int8Array<NonNullable, Buffer>,
341}
342
343impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
344    Clone for SparseUnionArray<T, VARIANTS, Buffer, OffsetItem>
345where
346    for<'a> i8: From<&'a T>,
347    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>: Clone,
348    Int8Array<NonNullable, Buffer>: Clone,
349{
350    fn clone(&self) -> Self {
351        Self {
352            variants: self.variants.clone(),
353            types: self.types.clone(),
354        }
355    }
356}
357
358impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
359    Default for SparseUnionArray<T, VARIANTS, Buffer, OffsetItem>
360where
361    for<'a> i8: From<&'a T>,
362    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>: Default,
363    Int8Array<NonNullable, Buffer>: Default,
364{
365    fn default() -> Self {
366        Self {
367            variants: Default::default(),
368            types: Int8Array::default(),
369        }
370    }
371}
372
373impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
374    Extend<T> for SparseUnionArray<T, VARIANTS, Buffer, OffsetItem>
375where
376    for<'a> i8: From<&'a T>,
377    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>: Extend<T>,
378    Int8Array<NonNullable, Buffer>: Extend<i8>,
379{
380    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
381        iter.into_iter().for_each(|item| {
382            self.types.extend(iter::once(i8::from(&item)));
383            self.variants.extend(iter::once(item));
384        });
385    }
386}
387
388impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
389    Length for SparseUnionArray<T, VARIANTS, Buffer, OffsetItem>
390where
391    for<'a> i8: From<&'a T>,
392{
393    fn len(&self) -> usize {
394        self.types.len()
395    }
396}
397
398impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
399    FromIterator<T> for SparseUnionArray<T, VARIANTS, Buffer, OffsetItem>
400where
401    for<'a> i8: From<&'a T>,
402    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>: Default + Extend<T>,
403    <Buffer as BufferType>::Buffer<i8>: Default + Extend<i8>,
404{
405    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
406        let (types, variants) = iter.into_iter().map(|item| (i8::from(&item), item)).unzip();
407        Self { variants, types }
408    }
409}
410
411/// Types that return a constructed `enum` by advancing
412/// iterator(s) of variants of a union array given the
413/// `type_id` for the variant.
414pub trait TypeIdIterator {
415    /// The rust enum
416    type Enum;
417
418    /// Advances the variant iterator(s) to get the next
419    /// value for the `type_id`.
420    fn next(&mut self, type_id: i8) -> Option<Self::Enum>;
421}
422
423/// Types that may be consumed to construct iterators of
424/// union array variants.
425pub trait UnionArrayIterators {
426    /// Type holding the variant iterators that may be
427    /// advanced to get the value for a `type_id`.
428    type VariantIterators: TypeIdIterator;
429
430    /// Constructs `VariantIterators` by consuming `Self`.
431    fn new_variant_iters(self) -> Self::VariantIterators;
432}
433
434/// Type holding the variant iterators of a union array
435type VarIters<T, const VARIANTS: usize, Buffer, OffsetItem, UnionLayout> = <<T as UnionArrayType<
436    VARIANTS,
437>>::Array<
438    Buffer,
439    OffsetItem,
440    UnionLayout,
441> as UnionArrayIterators>::VariantIterators;
442
443/// State required to iterate over union arrays
444pub struct UnionArrayIntoIter<
445    T: UnionArrayType<VARIANTS>,
446    const VARIANTS: usize,
447    Buffer: BufferType,
448    OffsetItem: Offset,
449    UnionLayout: UnionType,
450> where
451    for<'a> i8: From<&'a T>,
452    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, UnionLayout>: UnionArrayIterators,
453    <Buffer as BufferType>::Buffer<i8>: IntoIterator<Item = i8>,
454{
455    /// Type ids of the items in the union array
456    type_ids: <Int8Array<NonNullable, Buffer> as IntoIterator>::IntoIter,
457
458    /// Iterators of variants that may be advanced to get items for a type id
459    variant_iterators: VarIters<T, VARIANTS, Buffer, OffsetItem, UnionLayout>,
460}
461
462impl<
463    T: UnionArrayType<VARIANTS>,
464    const VARIANTS: usize,
465    Buffer: BufferType,
466    OffsetItem: Offset,
467    UnionLayout: UnionType,
468> Iterator for UnionArrayIntoIter<T, VARIANTS, Buffer, OffsetItem, UnionLayout>
469where
470    for<'a> i8: From<&'a T>,
471    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, UnionLayout>: UnionArrayIterators,
472    <Buffer as BufferType>::Buffer<i8>: IntoIterator<Item = i8>,
473    VarIters<T, VARIANTS, Buffer, OffsetItem, UnionLayout>: TypeIdIterator<Enum = T>,
474{
475    type Item = T;
476
477    fn next(&mut self) -> Option<Self::Item> {
478        self.type_ids.next().map(|type_id| {
479            self.variant_iterators
480                .next(type_id)
481                .expect("child arrays have correct length")
482        })
483    }
484}
485
486impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
487    IntoIterator for UnionArray<T, VARIANTS, SparseLayout, Buffer, OffsetItem>
488where
489    for<'a> i8: From<&'a T>,
490    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>: UnionArrayIterators,
491    <Buffer as BufferType>::Buffer<i8>: IntoIterator<Item = i8>,
492    VarIters<T, VARIANTS, Buffer, OffsetItem, SparseLayout>: TypeIdIterator<Enum = T>,
493{
494    type Item = T;
495    type IntoIter = UnionArrayIntoIter<T, VARIANTS, Buffer, OffsetItem, SparseLayout>;
496
497    fn into_iter(self) -> Self::IntoIter {
498        UnionArrayIntoIter {
499            variant_iterators: UnionArrayIterators::new_variant_iters(self.0.variants),
500            type_ids: self.0.types.into_iter(),
501        }
502    }
503}
504
505impl<T: UnionArrayType<VARIANTS>, const VARIANTS: usize, Buffer: BufferType, OffsetItem: Offset>
506    IntoIterator for UnionArray<T, VARIANTS, DenseLayout, Buffer, OffsetItem>
507where
508    for<'a> i8: From<&'a T>,
509    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>: UnionArrayIterators,
510    <Buffer as BufferType>::Buffer<i8>: IntoIterator<Item = i8>,
511    VarIters<T, VARIANTS, Buffer, OffsetItem, DenseLayout>: TypeIdIterator<Enum = T>,
512{
513    type Item = T;
514    type IntoIter = UnionArrayIntoIter<T, VARIANTS, Buffer, OffsetItem, DenseLayout>;
515
516    fn into_iter(self) -> Self::IntoIter {
517        UnionArrayIntoIter {
518            variant_iterators: UnionArrayIterators::new_variant_iters(self.0.variants),
519            type_ids: self.0.types.into_iter(),
520        }
521    }
522}
523
524#[cfg(test)]
525mod tests {
526    use super::*;
527    use crate::array::Uint32Array;
528    use std::marker::PhantomData;
529
530    #[test]
531    #[rustversion::attr(nightly, allow(non_local_definitions))]
532    #[allow(clippy::too_many_lines)]
533    fn simple() {
534        #[derive(Clone, Debug, PartialEq, Eq)]
535        enum Foo {
536            Bar(i32),
537            Baz(u32),
538        }
539
540        struct FooArray<Buffer: BufferType, UnionLayout: UnionType> {
541            bar: Int32Array<NonNullable, Buffer>,
542            baz: Uint32Array<NonNullable, Buffer>,
543            _ty: PhantomData<UnionLayout>, // we can also use a const generic instead?
544        }
545
546        impl<Buffer: BufferType> DenseOffset for FooArray<Buffer, DenseLayout> {
547            fn variant_len(&self, type_id: i8) -> usize {
548                match type_id {
549                    0 => self.bar.len(),
550                    1 => self.baz.len(),
551                    _ => panic!("bad type id"),
552                }
553            }
554        }
555
556        impl<Buffer: BufferType, UnionLayout: UnionType> Clone for FooArray<Buffer, UnionLayout>
557        where
558            Int32Array<NonNullable, Buffer>: Clone,
559            Uint32Array<NonNullable, Buffer>: Clone,
560        {
561            fn clone(&self) -> Self {
562                Self {
563                    bar: self.bar.clone(),
564                    baz: self.baz.clone(),
565                    _ty: self._ty,
566                }
567            }
568        }
569
570        impl<Buffer: BufferType, UnionLayout: UnionType> Default for FooArray<Buffer, UnionLayout>
571        where
572            Int32Array<NonNullable, Buffer>: Default,
573            Uint32Array<NonNullable, Buffer>: Default,
574        {
575            fn default() -> Self {
576                Self {
577                    bar: Int32Array::default(),
578                    baz: Uint32Array::default(),
579                    _ty: PhantomData,
580                }
581            }
582        }
583
584        impl<Buffer: BufferType> Extend<Foo> for FooArray<Buffer, DenseLayout>
585        where
586            Int32Array<NonNullable, Buffer>: Extend<i32>,
587            Uint32Array<NonNullable, Buffer>: Extend<u32>,
588        {
589            fn extend<T: IntoIterator<Item = Foo>>(&mut self, iter: T) {
590                iter.into_iter().for_each(|item| match item {
591                    Foo::Bar(x) => self.bar.extend(iter::once(x)),
592                    Foo::Baz(x) => self.baz.extend(iter::once(x)),
593                });
594            }
595        }
596
597        impl<Buffer: BufferType> Extend<Foo> for FooArray<Buffer, SparseLayout>
598        where
599            Int32Array<NonNullable, Buffer>: Extend<i32>,
600            Uint32Array<NonNullable, Buffer>: Extend<u32>,
601        {
602            fn extend<T: IntoIterator<Item = Foo>>(&mut self, iter: T) {
603                iter.into_iter().for_each(|item| match item {
604                    Foo::Bar(x) => {
605                        self.bar.extend(iter::once(x));
606                        self.baz.extend(iter::once(Default::default()));
607                    }
608                    Foo::Baz(x) => {
609                        self.baz.extend(iter::once(x));
610                        self.bar.extend(iter::once(Default::default()));
611                    }
612                });
613            }
614        }
615
616        struct FooArrayIntoIter<Buffer: BufferType, UnionLayout: UnionType>
617        where
618            Int32Array<NonNullable, Buffer>: IntoIterator,
619            Uint32Array<NonNullable, Buffer>: IntoIterator,
620        {
621            bar: <Int32Array<NonNullable, Buffer> as IntoIterator>::IntoIter,
622            baz: <Uint32Array<NonNullable, Buffer> as IntoIterator>::IntoIter,
623            _ty: PhantomData<UnionLayout>,
624        }
625
626        impl<Buffer: BufferType> TypeIdIterator for FooArrayIntoIter<Buffer, DenseLayout>
627        where
628            Int32Array<NonNullable, Buffer>: IntoIterator<Item = i32>,
629            Uint32Array<NonNullable, Buffer>: IntoIterator<Item = u32>,
630        {
631            type Enum = Foo;
632            fn next(&mut self, type_id: i8) -> Option<Self::Enum> {
633                match type_id {
634                    0 => self.bar.next().map(Foo::Bar),
635                    1 => self.baz.next().map(Foo::Baz),
636                    _ => panic!("type id greater than number of variants"),
637                }
638            }
639        }
640
641        impl<Buffer: BufferType> TypeIdIterator for FooArrayIntoIter<Buffer, SparseLayout>
642        where
643            Int32Array<NonNullable, Buffer>: IntoIterator<Item = i32>,
644            Uint32Array<NonNullable, Buffer>: IntoIterator<Item = u32>,
645        {
646            type Enum = Foo;
647            fn next(&mut self, type_id: i8) -> Option<Self::Enum> {
648                match type_id {
649                    0 => {
650                        self.baz.next();
651                        self.bar.next().map(Foo::Bar)
652                    }
653                    1 => {
654                        self.bar.next();
655                        self.baz.next().map(Foo::Baz)
656                    }
657                    _ => panic!("type id greater than number of variants"),
658                }
659            }
660        }
661
662        impl<Buffer: BufferType, UnionLayout: UnionType> UnionArrayIterators
663            for FooArray<Buffer, UnionLayout>
664        where
665            Int32Array<NonNullable, Buffer>: IntoIterator<Item = i32>,
666            Uint32Array<NonNullable, Buffer>: IntoIterator<Item = u32>,
667            FooArrayIntoIter<Buffer, UnionLayout>: TypeIdIterator,
668        {
669            type VariantIterators = FooArrayIntoIter<Buffer, UnionLayout>;
670
671            fn new_variant_iters(self) -> Self::VariantIterators {
672                Self::VariantIterators {
673                    bar: self.bar.into_iter(),
674                    baz: self.baz.into_iter(),
675                    _ty: PhantomData,
676                }
677            }
678        }
679
680        impl From<&Foo> for i8 {
681            fn from(value: &Foo) -> i8 {
682                match *value {
683                    Foo::Bar(_) => 0,
684                    Foo::Baz(_) => 1,
685                }
686            }
687        }
688
689        impl UnionArrayType<2> for Foo {
690            type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
691                FooArray<Buffer, UnionLayout>;
692        }
693
694        {
695            let input = vec![Foo::Bar(0), Foo::Baz(1), Foo::Baz(2), Foo::Bar(3)];
696            let mut dense_array = input
697                .clone()
698                .into_iter()
699                .collect::<UnionArray<Foo, { Foo::VARIANTS }>>();
700
701            assert_eq!(dense_array.0.types.0, [0, 1, 1, 0]);
702            assert_eq!(dense_array.0.offsets.0, [0, 0, 1, 1]);
703            assert_eq!(dense_array.0.variants.bar.0, [0, 3]);
704            assert_eq!(dense_array.0.variants.baz.0, [1, 2]);
705
706            assert_eq!(dense_array.clone().into_iter().collect::<Vec<_>>(), input);
707
708            dense_array.extend(iter::once(Foo::Bar(42)));
709            assert_eq!(dense_array.0.types.0, [0, 1, 1, 0, 0]);
710            assert_eq!(dense_array.0.offsets.0, [0, 0, 1, 1, 2]);
711            assert_eq!(dense_array.0.variants.bar.0, [0, 3, 42]);
712            assert_eq!(dense_array.0.variants.baz.0, [1, 2]);
713        };
714
715        {
716            let input = vec![Foo::Bar(-78), Foo::Baz(1), Foo::Baz(99)];
717            let mut sparse_array = input.clone().into_iter().collect::<UnionArray<
718                Foo,
719                { Foo::VARIANTS },
720                SparseLayout,
721            >>();
722
723            assert_eq!(sparse_array.0.types.0, [0, 1, 1]);
724            assert_eq!(
725                sparse_array.0.variants.bar.0,
726                [-78, i32::default(), i32::default()]
727            );
728            assert_eq!(sparse_array.0.variants.baz.0, [u32::default(), 1, 99]);
729
730            assert_eq!(sparse_array.clone().into_iter().collect::<Vec<_>>(), input);
731
732            sparse_array.extend(iter::once(Foo::Bar(42)));
733            assert_eq!(sparse_array.0.types.0, [0, 1, 1, 0]);
734            assert_eq!(
735                sparse_array.0.variants.bar.0,
736                [-78, i32::default(), i32::default(), 42]
737            );
738            assert_eq!(
739                sparse_array.0.variants.baz.0,
740                [u32::default(), 1, 99, u32::default()]
741            );
742        };
743    }
744
745    #[test]
746    #[cfg(feature = "derive")]
747    #[allow(clippy::too_many_lines)]
748    #[allow(clippy::type_complexity)]
749    #[allow(unused)]
750    #[rustversion::attr(nightly, allow(non_local_definitions))]
751    fn with_multiple_fields() {
752        use crate::{ArrayType, Length, array::ArrayTypeOf, offset};
753
754        #[derive(Clone, Debug, PartialEq, Eq)]
755        enum Foo {
756            Unit,
757            Unnamed(u8, u16),
758            Named { a: u32, b: u64, c: String },
759        }
760
761        impl From<&Foo> for i8 {
762            fn from(value: &Foo) -> i8 {
763                match *value {
764                    Foo::Unit => 0,
765                    Foo::Unnamed(..) => 1,
766                    Foo::Named { .. } => 2,
767                }
768            }
769        }
770
771        impl EnumVariant<0> for Foo {
772            type Data = ();
773
774            fn from_data(_value: Self::Data) -> Self {
775                Self::Unit
776            }
777        }
778
779        #[derive(ArrayType, Default)]
780        struct FooVariantUnnamed(u8, u16);
781
782        impl EnumVariant<1> for Foo {
783            type Data = FooVariantUnnamed;
784
785            fn from_data(value: Self::Data) -> Self {
786                Self::Unnamed(value.0, value.1)
787            }
788        }
789
790        #[derive(ArrayType, Default)]
791        struct FooVariantNamed {
792            a: u32,
793            b: u64,
794            c: String,
795        }
796
797        impl EnumVariant<2> for Foo {
798            type Data = FooVariantNamed;
799
800            fn from_data(value: Self::Data) -> Self {
801                Self::Named {
802                    a: value.a,
803                    b: value.b,
804                    c: value.c,
805                }
806            }
807        }
808
809        struct FooArray<Buffer: BufferType, UnionLayout: UnionType> {
810            unit: ArrayTypeOf<<Foo as EnumVariant<0>>::Data, Buffer, offset::NA, UnionLayout>,
811            unnamed: ArrayTypeOf<<Foo as EnumVariant<1>>::Data, Buffer, offset::NA, UnionLayout>,
812            named: ArrayTypeOf<<Foo as EnumVariant<2>>::Data, Buffer, offset::NA, UnionLayout>,
813        }
814
815        impl<Buffer: BufferType, UnionLayout: UnionType> Default for FooArray<Buffer, UnionLayout>
816        where
817            ArrayTypeOf<<Foo as EnumVariant<0>>::Data, Buffer, offset::NA, UnionLayout>: Default,
818            ArrayTypeOf<<Foo as EnumVariant<1>>::Data, Buffer, offset::NA, UnionLayout>: Default,
819            ArrayTypeOf<<Foo as EnumVariant<2>>::Data, Buffer, offset::NA, UnionLayout>: Default,
820        {
821            fn default() -> Self {
822                #[allow(clippy::default_trait_access)]
823                Self {
824                    unit: Default::default(),
825                    unnamed: Default::default(),
826                    named: Default::default(),
827                }
828            }
829        }
830
831        impl<Buffer: BufferType> Extend<Foo> for FooArray<Buffer, DenseLayout>
832        where
833            ArrayTypeOf<<Foo as EnumVariant<0>>::Data, Buffer, offset::NA, DenseLayout>:
834                Extend<<Foo as EnumVariant<0>>::Data>,
835            ArrayTypeOf<<Foo as EnumVariant<1>>::Data, Buffer, offset::NA, DenseLayout>:
836                Extend<<Foo as EnumVariant<1>>::Data>,
837            ArrayTypeOf<<Foo as EnumVariant<2>>::Data, Buffer, offset::NA, DenseLayout>:
838                Extend<<Foo as EnumVariant<2>>::Data>,
839        {
840            fn extend<T: IntoIterator<Item = Foo>>(&mut self, iter: T) {
841                iter.into_iter().for_each(|item| match item {
842                    Foo::Unit => self.unit.extend(iter::once(())),
843                    Foo::Unnamed(a, b) => self.unnamed.extend(iter::once(FooVariantUnnamed(a, b))),
844                    Foo::Named { a, b, c } => {
845                        self.named.extend(iter::once(FooVariantNamed { a, b, c }));
846                    }
847                });
848            }
849        }
850
851        impl<Buffer: BufferType> Extend<Foo> for FooArray<Buffer, SparseLayout>
852        where
853            ArrayTypeOf<<Foo as EnumVariant<0>>::Data, Buffer, offset::NA, SparseLayout>:
854                Extend<<Foo as EnumVariant<0>>::Data>,
855            ArrayTypeOf<<Foo as EnumVariant<1>>::Data, Buffer, offset::NA, SparseLayout>:
856                Extend<<Foo as EnumVariant<1>>::Data>,
857            ArrayTypeOf<<Foo as EnumVariant<2>>::Data, Buffer, offset::NA, SparseLayout>:
858                Extend<<Foo as EnumVariant<2>>::Data>,
859        {
860            fn extend<T: IntoIterator<Item = Foo>>(&mut self, iter: T) {
861                iter.into_iter().for_each(|item| match item {
862                    Foo::Unit => {
863                        self.unit.extend(iter::once(()));
864                        self.unnamed
865                            .extend(iter::once(FooVariantUnnamed::default()));
866                        self.named.extend(iter::once(FooVariantNamed::default()));
867                    }
868                    Foo::Unnamed(a, b) => {
869                        self.unit.extend(iter::once(()));
870                        self.unnamed.extend(iter::once(FooVariantUnnamed(a, b)));
871                        self.named.extend(iter::once(FooVariantNamed::default()));
872                    }
873                    Foo::Named { a, b, c } => {
874                        self.unit.extend(iter::once(()));
875                        self.unnamed
876                            .extend(iter::once(FooVariantUnnamed::default()));
877                        self.named.extend(iter::once(FooVariantNamed { a, b, c }));
878                    }
879                });
880            }
881        }
882
883        type FooEnumVariantArray<const INDEX: usize, Buffer, UnionLayout> =
884            ArrayTypeOf<<Foo as EnumVariant<INDEX>>::Data, Buffer, offset::NA, UnionLayout>;
885
886        struct FooArrayIntoIter<Buffer: BufferType, UnionLayout: UnionType>
887        where
888            FooEnumVariantArray<0, Buffer, UnionLayout>: IntoIterator,
889            FooEnumVariantArray<1, Buffer, UnionLayout>: IntoIterator,
890            FooEnumVariantArray<2, Buffer, UnionLayout>: IntoIterator,
891        {
892            unit: <FooEnumVariantArray<0, Buffer, UnionLayout> as IntoIterator>::IntoIter,
893            unnamed: <FooEnumVariantArray<1, Buffer, UnionLayout> as IntoIterator>::IntoIter,
894            named: <FooEnumVariantArray<2, Buffer, UnionLayout> as IntoIterator>::IntoIter,
895        }
896
897        impl<Buffer: BufferType> TypeIdIterator for FooArrayIntoIter<Buffer, DenseLayout>
898        where
899            FooEnumVariantArray<0, Buffer, DenseLayout>:
900                IntoIterator<Item = <Foo as EnumVariant<0>>::Data>,
901            FooEnumVariantArray<1, Buffer, DenseLayout>:
902                IntoIterator<Item = <Foo as EnumVariant<1>>::Data>,
903            FooEnumVariantArray<2, Buffer, DenseLayout>:
904                IntoIterator<Item = <Foo as EnumVariant<2>>::Data>,
905        {
906            type Enum = Foo;
907
908            fn next(&mut self, type_id: i8) -> Option<Self::Enum> {
909                match type_id {
910                    0 => self.unit.next().map(<Foo as EnumVariant<0>>::from_data),
911                    1 => self.unnamed.next().map(<Foo as EnumVariant<1>>::from_data),
912                    2 => self.named.next().map(<Foo as EnumVariant<2>>::from_data),
913                    _ => panic!("type id greater than number of variants"),
914                }
915            }
916        }
917
918        impl<Buffer: BufferType> TypeIdIterator for FooArrayIntoIter<Buffer, SparseLayout>
919        where
920            FooEnumVariantArray<0, Buffer, SparseLayout>:
921                IntoIterator<Item = <Foo as EnumVariant<0>>::Data>,
922            FooEnumVariantArray<1, Buffer, SparseLayout>:
923                IntoIterator<Item = <Foo as EnumVariant<1>>::Data>,
924            FooEnumVariantArray<2, Buffer, SparseLayout>:
925                IntoIterator<Item = <Foo as EnumVariant<2>>::Data>,
926        {
927            type Enum = Foo;
928
929            fn next(&mut self, type_id: i8) -> Option<Self::Enum> {
930                match type_id {
931                    0 => {
932                        let to_return = self.unit.next().map(<Foo as EnumVariant<0>>::from_data);
933                        self.unnamed.next();
934                        self.named.next();
935
936                        to_return
937                    }
938                    1 => {
939                        self.unit.next();
940                        let to_return = self.unnamed.next().map(<Foo as EnumVariant<1>>::from_data);
941                        self.named.next();
942
943                        to_return
944                    }
945                    2 => {
946                        self.unit.next();
947                        self.unnamed.next();
948                        let to_return = self.named.next().map(<Foo as EnumVariant<2>>::from_data);
949
950                        #[allow(clippy::let_and_return)]
951                        to_return
952                    }
953                    _ => panic!("type id greater than number of variants"),
954                }
955            }
956        }
957
958        impl<Buffer: BufferType, UnionLayout: UnionType> UnionArrayIterators
959            for FooArray<Buffer, UnionLayout>
960        where
961            FooEnumVariantArray<0, Buffer, SparseLayout>:
962                IntoIterator<Item = <Foo as EnumVariant<0>>::Data>,
963            FooEnumVariantArray<1, Buffer, SparseLayout>:
964                IntoIterator<Item = <Foo as EnumVariant<1>>::Data>,
965            FooEnumVariantArray<2, Buffer, SparseLayout>:
966                IntoIterator<Item = <Foo as EnumVariant<2>>::Data>,
967            FooArrayIntoIter<Buffer, UnionLayout>: TypeIdIterator,
968        {
969            type VariantIterators = FooArrayIntoIter<Buffer, UnionLayout>;
970
971            fn new_variant_iters(self) -> Self::VariantIterators {
972                Self::VariantIterators {
973                    unit: self.unit.into_iter(),
974                    unnamed: self.unnamed.into_iter(),
975                    named: self.named.into_iter(),
976                }
977            }
978        }
979
980        impl UnionArrayType<3> for Foo {
981            type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
982                FooArray<Buffer, UnionLayout>;
983        }
984
985        impl ArrayType<Foo> for Foo {
986            type Array<Buffer: BufferType, OffsetItem: offset::Offset, UnionLayout: UnionType> =
987                UnionArray<Foo, { Foo::VARIANTS }, UnionLayout>;
988        }
989
990        {
991            let input = vec![
992                Foo::Unit,
993                Foo::Unnamed(1, 2),
994                Foo::Named {
995                    a: 3,
996                    b: 4,
997                    c: "woo hoo!".to_owned(),
998                },
999            ];
1000            let dense_array = input
1001                .clone()
1002                .into_iter()
1003                .collect::<UnionArray<Foo, { Foo::VARIANTS }>>();
1004
1005            assert_eq!(dense_array.0.types.0, [0, 1, 2]);
1006            assert_eq!(dense_array.0.offsets.0, [0, 0, 0]);
1007            assert_eq!(dense_array.0.variants.unit.0.len(), 1);
1008            assert_eq!(dense_array.0.variants.unnamed.0.0.0, [1]);
1009            assert_eq!(dense_array.0.variants.unnamed.0.1.0, [2]);
1010            assert_eq!(dense_array.0.variants.named.0.a.0, [3]);
1011            assert_eq!(dense_array.0.variants.named.0.b.0, [4]);
1012
1013            assert_eq!(dense_array.into_iter().collect::<Vec<_>>(), input);
1014        };
1015
1016        {
1017            let input = vec![
1018                Foo::Unit,
1019                Foo::Unnamed(1, 2),
1020                Foo::Named {
1021                    a: 3,
1022                    b: 4,
1023                    c: "woo hoo!".to_owned(),
1024                },
1025            ];
1026            let sparse_array = input.clone().into_iter().collect::<UnionArray<
1027                Foo,
1028                { Foo::VARIANTS },
1029                SparseLayout,
1030            >>();
1031            assert_eq!(sparse_array.into_iter().collect::<Vec<_>>(), input);
1032        };
1033    }
1034
1035    #[test]
1036    #[cfg(feature = "derive")]
1037    #[rustversion::attr(nightly, allow(non_local_definitions))]
1038    fn derive() {
1039        use crate::ArrayType;
1040
1041        #[derive(ArrayType, Copy, Clone, Default, Debug, PartialEq, Eq)]
1042        enum Foo {
1043            #[default]
1044            Foo,
1045            Bar,
1046        }
1047
1048        #[derive(ArrayType, Clone, Copy, Debug, PartialEq, Eq)]
1049        enum Test {
1050            Foo { bar: u8 },
1051            Bar(bool),
1052            None,
1053        }
1054
1055        let foo_input = vec![Foo::Foo, Foo::Bar];
1056
1057        let dense_foo_array = foo_input.clone().into_iter().collect::<UnionArray<
1058            Foo,
1059            { Foo::VARIANTS },
1060            DenseLayout,
1061        >>();
1062        assert_eq!(dense_foo_array.len(), foo_input.len());
1063        let a = dense_foo_array.into_iter().collect::<Vec<_>>();
1064        assert_eq!(a, foo_input.clone());
1065
1066        let sparse_foo_array = foo_input.clone().into_iter().collect::<UnionArray<
1067            Foo,
1068            { <Foo as UnionArrayType<2>>::VARIANTS },
1069            SparseLayout,
1070        >>();
1071        assert_eq!(sparse_foo_array.len(), foo_input.len());
1072        assert_eq!(sparse_foo_array.into_iter().collect::<Vec<_>>(), foo_input);
1073
1074        let input = vec![
1075            Test::None,
1076            Test::Bar(true),
1077            Test::Foo { bar: 123 },
1078            Test::None,
1079        ];
1080        let mut dense_array = input
1081            .clone()
1082            .into_iter()
1083            .collect::<UnionArray<Test, { <Test as UnionArrayType<3>>::VARIANTS }>>();
1084        assert_eq!(dense_array.len(), 4);
1085        assert_eq!(dense_array.0.types.0, &[2, 1, 0, 2]);
1086        assert_eq!(dense_array.0.offsets.0, &[0, 0, 0, 1]);
1087        assert_eq!(dense_array.0.variants.0.0.bar.0, &[123]);
1088        assert_eq!(dense_array.0.variants.2.0.len(), 2);
1089        assert_eq!(
1090            dense_array.clone().into_iter().collect::<Vec<_>>(),
1091            input.clone()
1092        );
1093
1094        dense_array.extend(iter::once(Test::Foo { bar: 42 }));
1095        assert_eq!(dense_array.len(), 5);
1096        assert_eq!(dense_array.0.types.0, &[2, 1, 0, 2, 0]);
1097        assert_eq!(dense_array.0.offsets.0, &[0, 0, 0, 1, 1]);
1098        assert_eq!(dense_array.0.variants.0.0.bar.0, &[123, 42]);
1099        assert_eq!(dense_array.0.variants.2.0.len(), 2);
1100
1101        let sparse_array = input.clone().into_iter().collect::<UnionArray<
1102            Test,
1103            { <Test as UnionArrayType<3>>::VARIANTS },
1104            SparseLayout,
1105        >>();
1106        assert_eq!(sparse_array.len(), 4);
1107        assert_eq!(sparse_array.0.types.0, &[2, 1, 0, 2]);
1108        assert_eq!(sparse_array.0.variants.0.0.bar.0, &[0, 0, 123, 0]);
1109        assert_eq!(sparse_array.0.variants.2.0.len(), 4);
1110        assert_eq!(sparse_array.into_iter().collect::<Vec<_>>(), input);
1111    }
1112}
narrow/array/union.rs

narrow/array/
union.rs
