1use std::{
4 iter::{self, Map, Zip},
5 mem::{self, ManuallyDrop, MaybeUninit},
6};
7
8use crate::{
9 Index, Length,
10 bitmap::{Bitmap, BitmapRef, BitmapRefMut, ValidityBitmap},
11 buffer::{BufferMut, BufferType, VecBuffer},
12 nullability::{NonNullable, Nullability, Nullable},
13 validity::Validity,
14};
15
16use super::Array;
17
18pub struct FixedSizeListArray<
20 const N: usize,
21 T: Array,
22 Nullable: Nullability = NonNullable,
23 Buffer: BufferType = VecBuffer,
24>(pub(crate) Nullable::Collection<T, Buffer>);
25
26impl<const N: usize, T: Array + Length, Nullable: Nullability, Buffer: BufferType>
27 FixedSizeListArray<N, T, Nullable, Buffer>
28where
29 FixedSizeListArray<N, T, Nullable, Buffer>: Index + Length,
30{
31 pub fn iter(&self) -> FixedSizeListIter<'_, N, T, Nullable, Buffer> {
33 <&Self as IntoIterator>::into_iter(self)
34 }
35}
36
37impl<const N: usize, T: Array, Nullable: Nullability, Buffer: BufferType> Array
38 for FixedSizeListArray<N, T, Nullable, Buffer>
39{
40 type Item = Nullable::Item<[<T as Array>::Item; N]>;
41}
42
43impl<const N: usize, T: Array, Buffer: BufferType> BitmapRef
44 for FixedSizeListArray<N, T, Nullable, Buffer>
45{
46 type Buffer = Buffer;
47
48 fn bitmap_ref(&self) -> &Bitmap<Self::Buffer> {
49 self.0.bitmap_ref()
50 }
51}
52
53impl<const N: usize, T: Array, Buffer: BufferType> BitmapRefMut
54 for FixedSizeListArray<N, T, Nullable, Buffer>
55{
56 fn bitmap_ref_mut(&mut self) -> &mut Bitmap<Self::Buffer> {
57 self.0.bitmap_ref_mut()
58 }
59}
60
61impl<const N: usize, T: Array, Nullable: Nullability, Buffer: BufferType> Clone
62 for FixedSizeListArray<N, T, Nullable, Buffer>
63where
64 Nullable::Collection<T, Buffer>: Clone,
65{
66 fn clone(&self) -> Self {
67 Self(self.0.clone())
68 }
69}
70
71impl<const N: usize, T: Array, Nullable: Nullability, Buffer: BufferType> Default
72 for FixedSizeListArray<N, T, Nullable, Buffer>
73where
74 Nullable::Collection<T, Buffer>: Default,
75{
76 fn default() -> Self {
77 Self(Default::default())
78 }
79}
80
81impl<U, const N: usize, T: Array, Buffer: BufferType> Extend<[U; N]>
82 for FixedSizeListArray<N, T, NonNullable, Buffer>
83where
84 T: Extend<U>,
85{
86 fn extend<I: IntoIterator<Item = [U; N]>>(&mut self, iter: I) {
87 self.0.extend(iter.into_iter().flatten());
88 }
89}
90
91impl<U, const N: usize, T: Array, Buffer: BufferType> Extend<Option<[U; N]>>
92 for FixedSizeListArray<N, T, Nullable, Buffer>
93where
94 [U; N]: Default,
95 T: Extend<U>,
96 Bitmap<Buffer>: Extend<bool>,
97{
98 fn extend<I: IntoIterator<Item = Option<[U; N]>>>(&mut self, iter: I) {
99 self.0.data.extend(
100 iter.into_iter()
101 .inspect(|opt| {
102 self.0.validity.extend(iter::once(opt.is_some()));
103 })
104 .flat_map(Option::unwrap_or_default),
105 );
106 }
107}
108
109impl<const N: usize, T: Array, Buffer: BufferType>
110 From<FixedSizeListArray<N, T, NonNullable, Buffer>>
111 for FixedSizeListArray<N, T, Nullable, Buffer>
112where
113 T: Length,
114 Bitmap<Buffer>: FromIterator<bool>,
115{
116 fn from(value: FixedSizeListArray<N, T, NonNullable, Buffer>) -> Self {
117 Self(Validity::from(value.0))
118 }
119}
120
121impl<U, const N: usize, T: Array, Buffer: BufferType> FromIterator<[U; N]>
122 for FixedSizeListArray<N, T, NonNullable, Buffer>
123where
124 T: FromIterator<U>,
125{
126 fn from_iter<I: IntoIterator<Item = [U; N]>>(iter: I) -> Self {
127 Self(iter.into_iter().flatten().collect())
128 }
129}
130
131impl<U, const N: usize, T: Array, Buffer: BufferType> FromIterator<Option<[U; N]>>
132 for FixedSizeListArray<N, T, Nullable, Buffer>
133where
134 [U; N]: Default,
135 T: FromIterator<U>,
136 Buffer::Buffer<u8>: Default + BufferMut<u8> + Extend<u8>,
137{
138 fn from_iter<I: IntoIterator<Item = Option<[U; N]>>>(iter: I) -> Self {
139 let mut validity = Bitmap::default();
140 let data = iter
141 .into_iter()
142 .inspect(|opt| {
143 validity.extend(iter::once(opt.is_some()));
144 })
145 .flat_map(Option::unwrap_or_default)
146 .collect();
147 Self(Validity { data, validity })
148 }
149}
150
151impl<const N: usize, T: Array, Buffer: BufferType> Index
152 for FixedSizeListArray<N, T, NonNullable, Buffer>
153where
154 T: Index,
155{
156 type Item<'a>
157 = [<T as Index>::Item<'a>; N]
158 where
159 Self: 'a;
160
161 unsafe fn index_unchecked(&self, index: usize) -> Self::Item<'_> {
162 let mut data: [MaybeUninit<_>; N] = MaybeUninit::uninit().assume_init();
164 let start_index = index * N;
165 let end_index = start_index + N;
166 (start_index..end_index)
167 .enumerate()
168 .for_each(|(array_index, child_index)| {
169 data[array_index].write(self.0.index_unchecked(child_index));
170 });
171 mem::transmute_copy(&ManuallyDrop::new(data))
173 }
174}
175
176impl<const N: usize, T: Array, Buffer: BufferType> Index
177 for FixedSizeListArray<N, T, Nullable, Buffer>
178where
179 T: Index,
180{
181 type Item<'a>
182 = Option<[<T as Index>::Item<'a>; N]>
183 where
184 Self: 'a;
185
186 unsafe fn index_unchecked(&self, index: usize) -> Self::Item<'_> {
187 self.is_valid_unchecked(index).then(|| {
188 let mut data: [MaybeUninit<_>; N] = MaybeUninit::uninit().assume_init();
190 let start_index = index * N;
191 let end_index = start_index + N;
192 (start_index..end_index)
193 .enumerate()
194 .for_each(|(array_index, child_index)| {
195 data[array_index].write(self.0.data.index_unchecked(child_index));
197 });
198 mem::transmute_copy(&ManuallyDrop::new(data))
200 })
201 }
202}
203
204pub struct FixedSizeListIter<
206 'a,
207 const N: usize,
208 T: Array,
209 Nullable: Nullability,
210 Buffer: BufferType,
211> {
212 array: &'a FixedSizeListArray<N, T, Nullable, Buffer>,
214 index: usize,
216}
217
218impl<'a, const N: usize, T: Array, Nullable: Nullability, Buffer: BufferType> Iterator
219 for FixedSizeListIter<'a, N, T, Nullable, Buffer>
220where
221 FixedSizeListArray<N, T, Nullable, Buffer>: Length + Index,
222{
223 type Item = <FixedSizeListArray<N, T, Nullable, Buffer> as Index>::Item<'a>;
224
225 fn next(&mut self) -> Option<Self::Item> {
226 self.array
227 .index(self.index)
228 .into_iter()
229 .inspect(|_| {
230 self.index += 1;
231 })
232 .next()
233 }
234}
235
236impl<'a, const N: usize, T: Array, Nullable: Nullability, Buffer: BufferType> IntoIterator
237 for &'a FixedSizeListArray<N, T, Nullable, Buffer>
238where
239 FixedSizeListArray<N, T, Nullable, Buffer>: Index + Length,
240{
241 type Item = <FixedSizeListArray<N, T, Nullable, Buffer> as Index>::Item<'a>;
242 type IntoIter = FixedSizeListIter<'a, N, T, Nullable, Buffer>;
243
244 fn into_iter(self) -> Self::IntoIter {
245 FixedSizeListIter {
246 array: self,
247 index: 0,
248 }
249 }
250}
251
252pub struct FixedSizeArrayChunks<const N: usize, I: Iterator> {
254 iter: I,
256}
257
258impl<const N: usize, I: Iterator> FixedSizeArrayChunks<N, I> {
259 fn new(iter: I) -> Self {
261 Self { iter }
262 }
263}
264
265impl<const N: usize, I: Iterator> Iterator for FixedSizeArrayChunks<N, I> {
266 type Item = [I::Item; N];
267
268 fn next(&mut self) -> Option<Self::Item> {
269 let mut data: [MaybeUninit<I::Item>; N] =
270 unsafe { MaybeUninit::uninit().assume_init() };
273
274 let mut total_elements_written: usize = 0;
275 self.iter
276 .by_ref()
277 .take(N)
278 .enumerate()
279 .for_each(|(array_index, val)| {
280 data[array_index].write(val);
281 total_elements_written += 1;
282 });
283
284 assert!(total_elements_written <= N);
285
286 if total_elements_written == N {
287 Some(data.map(|elem| {
288 unsafe { elem.assume_init() }
291 }))
292 } else {
293 for elem in &mut data[0..total_elements_written] {
295 unsafe {
298 elem.assume_init_drop();
300 }
301 }
302 None
303 }
304 }
305}
306
307impl<const N: usize, T: Array, Buffer: BufferType> IntoIterator
308 for FixedSizeListArray<N, T, NonNullable, Buffer>
309where
310 T: IntoIterator,
311 FixedSizeArrayChunks<N, <T as IntoIterator>::IntoIter>:
312 IntoIterator<Item = [<T as IntoIterator>::Item; N]>,
313{
314 type Item = [<T as IntoIterator>::Item; N];
315 type IntoIter = FixedSizeArrayChunks<N, <T as IntoIterator>::IntoIter>;
316
317 fn into_iter(self) -> Self::IntoIter {
318 FixedSizeArrayChunks::<N, _>::new(self.0.into_iter())
319 }
320}
321
322impl<const N: usize, T: Array, Buffer: BufferType> IntoIterator
323 for FixedSizeListArray<N, T, Nullable, Buffer>
324where
325 T: IntoIterator,
326 Bitmap<Buffer>: IntoIterator<Item = bool>,
327 FixedSizeArrayChunks<N, <T as IntoIterator>::IntoIter>:
328 IntoIterator<Item = [<T as IntoIterator>::Item; N]>,
329{
330 type Item = Option<[<T as IntoIterator>::Item; N]>;
331 type IntoIter = Map<
332 Zip<
333 <Bitmap<Buffer> as IntoIterator>::IntoIter,
334 <FixedSizeArrayChunks<N, <T as IntoIterator>::IntoIter> as IntoIterator>::IntoIter,
335 >,
336 fn((bool, [<T as IntoIterator>::Item; N])) -> Self::Item,
337 >;
338
339 fn into_iter(self) -> Self::IntoIter {
340 self.0
341 .validity
342 .into_iter()
343 .zip(FixedSizeArrayChunks::<N, _>::new(self.0.data.into_iter()))
344 .map(|(validity, value)| validity.then_some(value))
345 }
346}
347
348impl<const N: usize, T: Array, Nullable: Nullability, Buffer: BufferType> Length
349 for FixedSizeListArray<N, T, Nullable, Buffer>
350where
351 Nullable::Collection<T, Buffer>: Length,
352{
353 fn len(&self) -> usize {
354 if Nullable::NULLABLE {
355 self.0.len()
357 } else {
358 self.0.len() / N
359 }
360 }
361}
362
363impl<const N: usize, T: Array, Buffer: BufferType> ValidityBitmap
364 for FixedSizeListArray<N, T, Nullable, Buffer>
365{
366}
367
368#[cfg(test)]
369mod tests {
370 use crate::array::{FixedSizePrimitiveArray, StringArray};
371
372 use super::*;
373
374 #[test]
375 fn from_iter() {
376 {
377 let input_non_nullable = [[1_u8, 2], [3, 4]];
378 let array_non_nullable = input_non_nullable
379 .into_iter()
380 .collect::<FixedSizeListArray<2, FixedSizePrimitiveArray<u8>>>();
381 assert_eq!(array_non_nullable.len(), 2);
382 };
383
384 {
385 let input_inner_nullable = [[Some(1_u8), None], [Some(3), None]];
386 let array_inner_nullable = input_inner_nullable
387 .into_iter()
388 .collect::<FixedSizeListArray<2, FixedSizePrimitiveArray<u8, Nullable>, NonNullable>>();
389 assert_eq!(array_inner_nullable.len(), 2);
390 assert_eq!(
391 array_inner_nullable.into_iter().collect::<Vec<_>>(),
392 input_inner_nullable
393 );
394 };
395
396 {
397 let input_outer_nullable = [Some([1_u8, 1_u8]), Some([1_u8, 1_u8]), None];
398 let array_outer_nullable = input_outer_nullable
399 .into_iter()
400 .collect::<FixedSizeListArray<2, FixedSizePrimitiveArray<u8, NonNullable>, Nullable>>();
401 assert_eq!(array_outer_nullable.len(), 3);
402 assert_eq!(
403 array_outer_nullable.into_iter().collect::<Vec<_>>(),
404 input_outer_nullable
405 );
406 };
407
408 {
409 let input_both_nullable = [Some([Some(1_u8), None]), None];
410 let array_both_nullable = input_both_nullable
411 .into_iter()
412 .collect::<FixedSizeListArray<2, FixedSizePrimitiveArray<u8, Nullable>, Nullable>>(
413 );
414 assert_eq!(array_both_nullable.len(), 2);
415 assert_eq!(
416 array_both_nullable.into_iter().collect::<Vec<_>>(),
417 input_both_nullable
418 );
419 };
420
421 {
422 let input_nested_innermost_nullable = [
423 [
424 [Some(1_u8), None, Some(1_u8)],
425 [Some(3_u8), None, Some(1_u8)],
426 ],
427 [
428 [Some(2_u8), None, Some(1_u8)],
429 [Some(5_u8), None, Some(1_u8)],
430 ],
431 ];
432 let array_nested_innermost_nullable = input_nested_innermost_nullable
433 .into_iter()
434 .collect::<FixedSizeListArray<
435 2,
436 FixedSizeListArray<3, FixedSizePrimitiveArray<u8, Nullable>, NonNullable>,
437 NonNullable,
438 >>();
439 assert_eq!(array_nested_innermost_nullable.len(), 2);
440 assert_eq!(
441 array_nested_innermost_nullable
442 .into_iter()
443 .collect::<Vec<_>>(),
444 input_nested_innermost_nullable
445 );
446 };
447
448 {
449 let input_nested_all_nullable = [
450 None,
451 Some([
452 None,
453 Some([Some(1_u8), None, Some(2)]),
454 Some([Some(3), None, Some(1)]),
455 None,
456 ]),
457 Some([
458 Some([Some(2), None, Some(1)]),
459 None,
460 None,
461 Some([Some(5), None, Some(6)]),
462 ]),
463 ];
464 let array_nested_all_nullable = input_nested_all_nullable
465 .into_iter()
466 .collect::<FixedSizeListArray<
467 4,
468 FixedSizeListArray<3, FixedSizePrimitiveArray<u8, Nullable>, Nullable>,
469 Nullable,
470 >>();
471 assert_eq!(array_nested_all_nullable.len(), 3);
472 assert_eq!(
473 array_nested_all_nullable.into_iter().collect::<Vec<_>>(),
474 input_nested_all_nullable
475 );
476 };
477 }
478
479 #[test]
480 fn from_iter_variable_size() {
481 {
482 let input_string_non_nullable = [
483 ["hello".to_owned(), "world".to_owned()],
484 ["!".to_owned(), "!".to_owned()],
485 ];
486 let array_string_non_nullable = input_string_non_nullable
487 .clone()
488 .into_iter()
489 .collect::<FixedSizeListArray<2, StringArray>>();
490 assert_eq!(array_string_non_nullable.len(), 2);
491 assert_eq!(
492 array_string_non_nullable.into_iter().collect::<Vec<_>>(),
493 input_string_non_nullable
494 );
495 };
496
497 {
498 let input_string_nested_all_nullable = [
499 None,
500 Some([
501 Some([Some("hello".to_owned()), None, Some("from".to_owned())]),
502 Some([Some("the".to_owned()), None, Some("other".to_owned())]),
503 None,
504 None,
505 ]),
506 Some([
507 None,
508 Some([Some("side".to_owned()), None, Some("hello".to_owned())]),
509 None,
510 Some([Some("its".to_owned()), None, Some("me".to_owned())]),
511 ]),
512 ];
513 let array_string_nested_all_nullable = input_string_nested_all_nullable
514 .clone()
515 .into_iter()
516 .collect::<FixedSizeListArray<
517 4,
518 FixedSizeListArray<3, StringArray<Nullable>, Nullable>,
519 Nullable,
520 >>();
521 assert_eq!(array_string_nested_all_nullable.len(), 3);
522 assert_eq!(
523 array_string_nested_all_nullable
524 .into_iter()
525 .collect::<Vec<_>>(),
526 input_string_nested_all_nullable
527 );
528 };
529
530 {
531 let input_string_even_more_nested = [
532 Some([
533 Some([Some(["hello".to_owned()]), None, Some(["from".to_owned()])]),
534 Some([Some(["the".to_owned()]), None, Some(["other".to_owned()])]),
535 None,
536 None,
537 ]),
538 None,
539 Some([
540 None,
541 Some([Some(["side".to_owned()]), None, Some(["hello".to_owned()])]),
542 None,
543 Some([Some(["its".to_owned()]), None, Some(["me".to_owned()])]),
544 ]),
545 ];
546 let array_string_even_more_nested = input_string_even_more_nested
547 .clone()
548 .into_iter()
549 .collect::<FixedSizeListArray<
550 4,
551 FixedSizeListArray<
552 3,
553 FixedSizeListArray<1, StringArray<NonNullable>, Nullable>,
554 Nullable,
555 >,
556 Nullable,
557 >>();
558 assert_eq!(array_string_even_more_nested.len(), 3);
559 assert_eq!(
560 array_string_even_more_nested
561 .into_iter()
562 .collect::<Vec<_>>(),
563 input_string_even_more_nested
564 );
565 };
566 }
567
568 #[test]
569 fn index() {
570 let input = [[1_u8, 2], [3, 4]];
571 let array = input
572 .into_iter()
573 .collect::<FixedSizeListArray<2, FixedSizePrimitiveArray<u8>>>();
574 assert_eq!(array.index(0), Some([&1, &2]));
575 assert_eq!(array.index(1), Some([&3, &4]));
576
577 let input_string = [["hello", "world"], ["!", "!"]];
578 let array_string = input_string
579 .into_iter()
580 .collect::<FixedSizeListArray<2, StringArray>>();
581 assert_eq!(array_string.index(0), Some(["hello", "world"]));
582 assert_eq!(array_string.index(1), Some(["!", "!"]));
583
584 let input_nullable_string = [Some(["hello", "world"]), None];
585 let array_nullable_string = input_nullable_string
586 .into_iter()
587 .collect::<FixedSizeListArray<2, StringArray, Nullable>>();
588 assert_eq!(
589 array_nullable_string.index(0),
590 Some(Some(["hello", "world"]))
591 );
592 assert_eq!(array_nullable_string.index(1), Some(None));
593 assert_eq!(array_nullable_string.index(2), None);
594
595 let input_nullable_string_nullable = [
596 Some([Some("hello"), None]),
597 None,
598 Some([None, Some("world")]),
599 ];
600 let array_nullable_string_nullable = input_nullable_string_nullable
601 .into_iter()
602 .collect::<FixedSizeListArray<2, StringArray<Nullable>, Nullable>>(
603 );
604 assert_eq!(
605 array_nullable_string_nullable.index(0),
606 Some(Some([Some("hello"), None]))
607 );
608 assert_eq!(array_nullable_string_nullable.index(1), Some(None));
609 assert_eq!(
610 array_nullable_string_nullable.index(2),
611 Some(Some([None, Some("world")]))
612 );
613 assert_eq!(array_nullable_string_nullable.index(3), None);
614 }
615
616 #[test]
617 fn fixed_size_array_chunks() {
618 {
619 let input = vec![0, 1, 2, 3, 4, 5, 6, 7, 8];
620 let array_chunks = FixedSizeArrayChunks::<3, _>::new(input.into_iter());
621 assert_eq!(
622 array_chunks.into_iter().collect::<Vec<_>>(),
623 vec![[0, 1, 2], [3, 4, 5], [6, 7, 8]]
624 );
625 };
626
627 {
628 let input = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
630 let array_chunks = FixedSizeArrayChunks::<3, _>::new(input.into_iter());
631 assert_eq!(
632 array_chunks.into_iter().collect::<Vec<_>>(),
633 vec![[0, 1, 2], [3, 4, 5], [6, 7, 8]]
634 );
635 }
636 }
637
638 #[test]
639 fn into_iter() {
640 let input = [[1_u8, 2], [3, 4]];
641 let array = input
642 .into_iter()
643 .collect::<FixedSizeListArray<2, FixedSizePrimitiveArray<u8>>>();
644 assert_eq!(array.into_iter().collect::<Vec<_>>(), [[1, 2], [3, 4]]);
645
646 let input_string = [["hello", "world"], ["!", "!"]];
647 let array_string = input_string
648 .into_iter()
649 .collect::<FixedSizeListArray<2, StringArray>>();
650 assert_eq!(array_string.into_iter().collect::<Vec<_>>(), input_string);
651
652 let input_nullable_string = [
653 Some(["hello".to_owned(), "world".to_owned()]),
654 None,
655 Some(["hello".to_owned(), "again".to_owned()]),
656 ];
657 let array_nullable_string = input_nullable_string
658 .clone()
659 .into_iter()
660 .collect::<FixedSizeListArray<2, StringArray, Nullable>>();
661 assert_eq!(
662 array_nullable_string.into_iter().collect::<Vec<_>>(),
663 input_nullable_string
664 );
665
666 let input_nullable_string_nullable = [
667 Some([Some("hello".to_owned()), None]),
668 None,
669 Some([None, Some("world".to_owned())]),
670 None,
671 Some([Some("hello".to_owned()), Some("again".to_owned())]),
672 ];
673 let array_nullable_string_nullable = input_nullable_string_nullable
674 .clone()
675 .into_iter()
676 .collect::<FixedSizeListArray<2, StringArray<Nullable>, Nullable>>(
677 );
678 assert_eq!(
679 array_nullable_string_nullable
680 .into_iter()
681 .collect::<Vec<_>>(),
682 input_nullable_string_nullable
683 );
684
685 let input_nested = [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 0], [0, 0]]];
686 let array_nested = input_nested
687 .into_iter()
688 .collect::<FixedSizeListArray<3, FixedSizeListArray<2, FixedSizePrimitiveArray<u8>>>>();
689
690 assert_eq!(array_nested.0.0.0, [1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 0]);
691 assert_eq!(
692 array_nested.into_iter().collect::<Vec<_>>(),
693 [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 0], [0, 0]]]
694 );
695 }
696}