narrow/array/
variable_size_list.rs

1//!Array with variable-size list elements.
2
3use crate::{
4    Index, Length,
5    array::Array,
6    bitmap::{Bitmap, BitmapRef, BitmapRefMut, ValidityBitmap},
7    buffer::{BufferType, VecBuffer},
8    nullability::{NonNullable, Nullability, Nullable},
9    offset::{Offset, Offsets},
10};
11use std::fmt::{Debug, Formatter, Result};
12
13/// Array with variable-size list elements.
14pub struct VariableSizeListArray<
15    T: Array, // todo(mbrobbel): move this bound?
16    Nullable: Nullability = NonNullable,
17    OffsetItem: Offset = i32,
18    Buffer: BufferType = VecBuffer,
19>(pub Offsets<T, Nullable, OffsetItem, Buffer>);
20
21impl<T: Array, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> Array
22    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
23{
24    type Item = Nullable::Item<Vec<T>>;
25}
26
27impl<T: Array, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> Debug
28    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
29where
30    Offsets<T, Nullable, OffsetItem, Buffer>: Debug,
31{
32    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
33        f.debug_tuple("VariableSizeListArray")
34            .field(&self.0)
35            .finish()
36    }
37}
38
39impl<T: Array, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> Clone
40    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
41where
42    Offsets<T, Nullable, OffsetItem, Buffer>: Clone,
43{
44    fn clone(&self) -> Self {
45        Self(self.0.clone())
46    }
47}
48
49impl<T: Array, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> Default
50    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
51where
52    Offsets<T, Nullable, OffsetItem, Buffer>: Default,
53{
54    fn default() -> Self {
55        Self(Offsets::default())
56    }
57}
58
59impl<T: Array, U, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> Extend<U>
60    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
61where
62    Offsets<T, Nullable, OffsetItem, Buffer>: Extend<U>,
63{
64    fn extend<I: IntoIterator<Item = U>>(&mut self, iter: I) {
65        self.0.extend(iter);
66    }
67}
68
69impl<T: Array, OffsetItem: Offset, Buffer: BufferType>
70    From<VariableSizeListArray<T, NonNullable, OffsetItem, Buffer>>
71    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
72where
73    Offsets<T, NonNullable, OffsetItem, Buffer>: Into<Offsets<T, Nullable, OffsetItem, Buffer>>,
74{
75    fn from(value: VariableSizeListArray<T, NonNullable, OffsetItem, Buffer>) -> Self {
76        Self(value.0.into())
77    }
78}
79
80impl<T: Array, U, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> FromIterator<U>
81    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
82where
83    Offsets<T, Nullable, OffsetItem, Buffer>: FromIterator<U>,
84{
85    fn from_iter<I: IntoIterator<Item = U>>(iter: I) -> Self {
86        Self(iter.into_iter().collect())
87    }
88}
89
90impl<T: Array, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> IntoIterator
91    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
92where
93    Offsets<T, Nullable, OffsetItem, Buffer>: IntoIterator,
94{
95    type IntoIter = <Offsets<T, Nullable, OffsetItem, Buffer> as IntoIterator>::IntoIter;
96    type Item = <Offsets<T, Nullable, OffsetItem, Buffer> as IntoIterator>::Item;
97
98    fn into_iter(self) -> Self::IntoIter {
99        self.0.into_iter()
100    }
101}
102
103impl<T: Array, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> Index
104    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
105where
106    Offsets<T, Nullable, OffsetItem, Buffer>: Index,
107{
108    type Item<'a>
109        = <Offsets<T, Nullable, OffsetItem, Buffer> as Index>::Item<'a>
110    where
111        Self: 'a;
112
113    unsafe fn index_unchecked(&self, index: usize) -> Self::Item<'_> {
114        self.0.index_unchecked(index)
115    }
116}
117
118impl<T: Array, Nullable: Nullability, OffsetItem: Offset, Buffer: BufferType> Length
119    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
120where
121    Offsets<T, Nullable, OffsetItem, Buffer>: Length,
122{
123    fn len(&self) -> usize {
124        self.0.len()
125    }
126}
127
128impl<T: Array, OffsetItem: Offset, Buffer: BufferType> BitmapRef
129    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
130{
131    type Buffer = Buffer;
132
133    fn bitmap_ref(&self) -> &Bitmap<Self::Buffer> {
134        self.0.bitmap_ref()
135    }
136}
137
138impl<T: Array, OffsetItem: Offset, Buffer: BufferType> BitmapRefMut
139    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
140{
141    fn bitmap_ref_mut(&mut self) -> &mut Bitmap<Self::Buffer> {
142        self.0.bitmap_ref_mut()
143    }
144}
145
146impl<T: Array, OffsetItem: Offset, Buffer: BufferType> ValidityBitmap
147    for VariableSizeListArray<T, Nullable, OffsetItem, Buffer>
148{
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::array::FixedSizePrimitiveArray;
155
156    #[test]
157    fn from_iter() {
158        let input = vec![vec![1], vec![2, 3], vec![4]];
159        let array = input
160            .into_iter()
161            .collect::<VariableSizeListArray<FixedSizePrimitiveArray<u8>>>();
162        assert_eq!(array.0.data.0, &[1, 2, 3, 4]);
163        assert_eq!(array.0.offsets, &[0, 1, 3, 4]);
164    }
165
166    #[test]
167    fn from_iter_nullable() {
168        let input = vec![Some(vec![1]), None, Some(vec![2, 3]), Some(vec![4])];
169        let array = input
170            .into_iter()
171            .collect::<VariableSizeListArray<FixedSizePrimitiveArray<u8>, Nullable>>();
172        assert_eq!(array.len(), 4);
173        assert_eq!(array.0.data.0, &[1, 2, 3, 4]);
174        assert_eq!(array.0.offsets.bitmap_ref().is_valid(0), Some(true));
175        assert_eq!(array.0.offsets.bitmap_ref().is_null(1), Some(true));
176        assert_eq!(array.0.offsets.bitmap_ref().is_valid(2), Some(true));
177        assert_eq!(array.0.offsets.bitmap_ref().is_valid(3), Some(true));
178        assert_eq!(array.0.offsets.as_ref().as_slice(), &[0, 1, 1, 3, 4]);
179    }
180
181    #[test]
182    fn from_iter_nested() {
183        let input = vec![vec![vec![1, 2, 3], vec![1, 2, 3]], vec![vec![4, 5, 6]]];
184        let array = input
185            .into_iter()
186            .collect::<VariableSizeListArray<VariableSizeListArray<FixedSizePrimitiveArray<u8>>>>();
187        assert_eq!(array.0.data.0.data.0, &[1, 2, 3, 1, 2, 3, 4, 5, 6]);
188        assert_eq!(array.0.offsets, &[0, 2, 3]);
189        assert_eq!(array.0.data.0.offsets, &[0, 3, 6, 9]);
190
191        let input_3 = vec![vec![
192            vec![vec![1, 2, 3], vec![1, 2, 3]],
193            vec![vec![4, 5, 6]],
194        ]];
195        let array_3 = input_3.into_iter().collect::<VariableSizeListArray<
196            VariableSizeListArray<VariableSizeListArray<FixedSizePrimitiveArray<u8>>>,
197        >>();
198        assert_eq!(array_3.0.data.0.data.0.data.0, &[1, 2, 3, 1, 2, 3, 4, 5, 6]);
199        assert_eq!(array_3.0.offsets, &[0, 2]);
200        assert_eq!(array_3.0.data.0.offsets, &[0, 2, 3]);
201        assert_eq!(array_3.0.data.0.data.0.offsets, &[0, 3, 6, 9]);
202    }
203
204    #[test]
205    fn from_iter_nested_nullable() {
206        let input = vec![
207            None,
208            Some(vec![
209                None,
210                Some(vec![None, Some(vec![1, 2, 3]), Some(vec![1, 2, 3])]),
211                Some(vec![None, Some(vec![4, 5, 6])]),
212            ]),
213        ];
214        let array = input.into_iter().collect::<VariableSizeListArray<
215            VariableSizeListArray<
216                VariableSizeListArray<FixedSizePrimitiveArray<u8>, Nullable>,
217                Nullable,
218            >,
219            Nullable,
220        >>();
221        assert_eq!(array.0.data.0.data.0.data.0, &[1, 2, 3, 1, 2, 3, 4, 5, 6]);
222        assert_eq!(array.0.offsets.as_ref(), &[0, 0, 3]);
223        assert_eq!(array.0.data.0.offsets.as_ref(), &[0, 0, 3, 5]);
224        assert_eq!(array.0.data.0.data.0.offsets.as_ref(), &[0, 0, 3, 6, 6, 9]);
225        assert_eq!(array.is_null(0), Some(true));
226        assert_eq!(array.is_valid(1), Some(true));
227        assert_eq!(array.0.data.is_null(0), Some(true));
228        assert_eq!(array.0.data.is_valid(1), Some(true));
229        assert_eq!(array.0.data.0.data.is_null(0), Some(true));
230        assert_eq!(array.0.data.0.data.is_valid(1), Some(true));
231        assert_eq!(array.0.data.0.data.0.is_null(0), Some(true));
232        assert_eq!(array.0.data.0.data.0.is_valid(1), Some(true));
233
234        let input_3 = vec![
235            None,
236            Some(vec![
237                None,
238                Some(vec![
239                    None,
240                    Some(vec![None, Some(2), Some(3)]),
241                    Some(vec![None, Some(2), Some(3)]),
242                ]),
243                Some(vec![None, Some(vec![None, Some(5), Some(6)])]),
244            ]),
245        ];
246        let array_3 = input_3.into_iter().collect::<VariableSizeListArray<
247            VariableSizeListArray<
248                VariableSizeListArray<FixedSizePrimitiveArray<u8, Nullable>, Nullable>,
249                Nullable,
250            >,
251            Nullable,
252        >>();
253        assert_eq!(
254            array_3.0.data.0.data.0.data.0.as_ref(),
255            &[
256                u8::default(),
257                2,
258                3,
259                u8::default(),
260                2,
261                3,
262                u8::default(),
263                5,
264                6
265            ]
266        );
267        assert_eq!(array_3.0.offsets.as_ref(), &[0, 0, 3]);
268        assert_eq!(array_3.0.data.0.offsets.as_ref(), &[0, 0, 3, 5]);
269        assert_eq!(
270            array_3.0.data.0.data.0.offsets.as_ref(),
271            &[0, 0, 3, 6, 6, 9]
272        );
273        assert_eq!(array_3.is_null(0), Some(true));
274        assert_eq!(array_3.is_valid(1), Some(true));
275        assert_eq!(array_3.0.data.is_null(0), Some(true));
276        assert_eq!(array_3.0.data.is_valid(1), Some(true));
277        assert_eq!(array_3.0.data.0.data.is_null(0), Some(true));
278        assert_eq!(array_3.0.data.0.data.is_valid(1), Some(true));
279        assert_eq!(array_3.0.data.0.data.0.is_null(0), Some(true));
280        assert_eq!(array_3.0.data.0.data.0.is_valid(1), Some(true));
281        assert_eq!(array_3.0.data.0.data.0.data.0.is_null(0), Some(true));
282        assert_eq!(array_3.0.data.0.data.0.data.0.is_valid(1), Some(true));
283    }
284
285    #[test]
286    fn index() {
287        let input = vec![vec![vec![1, 2, 3], vec![1, 2, 3]], vec![vec![4, 5, 6]]];
288        let array = input
289            .into_iter()
290            .collect::<VariableSizeListArray<VariableSizeListArray<FixedSizePrimitiveArray<u8>>>>();
291        assert_eq!(
292            array
293                .index_checked(0)
294                .flatten()
295                .copied()
296                .collect::<Vec<_>>(),
297            [1, 2, 3, 1, 2, 3]
298        );
299        assert_eq!(
300            array
301                .index_checked(1)
302                .next()
303                .expect("a value")
304                .copied()
305                .collect::<Vec<_>>(),
306            [4, 5, 6]
307        );
308        assert!(array.index(2).is_none());
309    }
310
311    #[test]
312    fn into_iter_nested() {
313        let input = vec![
314            vec![vec![1, 2, 3], vec![1, 2, 3]],
315            vec![],
316            vec![vec![], vec![]],
317            vec![vec![4, 5, 6]],
318        ];
319        let array = input
320            .clone()
321            .into_iter()
322            .collect::<VariableSizeListArray<VariableSizeListArray<FixedSizePrimitiveArray<u8>>>>();
323        assert_eq!(array.into_iter().collect::<Vec<_>>(), input);
324    }
325
326    #[test]
327    fn into_iter_nested_nullable() {
328        let input = vec![
329            None,
330            Some(vec![]),
331            Some(vec![None, Some(vec![])]),
332            Some(vec![
333                None,
334                Some(vec![None, Some(vec![1, 2, 3]), Some(vec![1, 2, 3])]),
335                Some(vec![None, Some(vec![4, 5, 6])]),
336            ]),
337        ];
338        let array = input.clone().into_iter().collect::<VariableSizeListArray<
339            VariableSizeListArray<
340                VariableSizeListArray<FixedSizePrimitiveArray<u8>, Nullable>,
341                Nullable,
342            >,
343            Nullable,
344        >>();
345        assert_eq!(array.into_iter().collect::<Vec<_>>(), input);
346    }
347}
narrow/array/variable_size_list.rs

narrow/array/
variable_size_list.rs
