1use 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
12pub trait StructArrayType: ArrayType<Self> {
14 type Array<Buffer: BufferType>; }
20
21pub 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 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 #[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 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] );
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 };
537 }
538}