1use crate::builder::ArrayBuilder;
19use crate::{ArrayRef, GenericListViewArray, OffsetSizeTrait};
20use arrow_buffer::{Buffer, NullBufferBuilder, ScalarBuffer};
21use arrow_schema::{Field, FieldRef};
22use std::any::Any;
23use std::sync::Arc;
24
25#[derive(Debug)]
27pub struct GenericListViewBuilder<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> {
28 offsets_builder: Vec<OffsetSize>,
29 sizes_builder: Vec<OffsetSize>,
30 null_buffer_builder: NullBufferBuilder,
31 values_builder: T,
32 field: Option<FieldRef>,
33 current_offset: OffsetSize,
34}
35
36impl<O: OffsetSizeTrait, T: ArrayBuilder + Default> Default for GenericListViewBuilder<O, T> {
37 fn default() -> Self {
38 Self::new(T::default())
39 }
40}
41
42impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> ArrayBuilder
43 for GenericListViewBuilder<OffsetSize, T>
44{
45 fn as_any(&self) -> &dyn Any {
47 self
48 }
49
50 fn as_any_mut(&mut self) -> &mut dyn Any {
52 self
53 }
54
55 fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
57 self
58 }
59
60 fn len(&self) -> usize {
62 self.null_buffer_builder.len()
63 }
64
65 fn finish(&mut self) -> ArrayRef {
67 Arc::new(self.finish())
68 }
69
70 fn finish_cloned(&self) -> ArrayRef {
72 Arc::new(self.finish_cloned())
73 }
74}
75
76impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListViewBuilder<OffsetSize, T> {
77 pub fn new(values_builder: T) -> Self {
79 let capacity = values_builder.len();
80 Self::with_capacity(values_builder, capacity)
81 }
82
83 pub fn with_capacity(values_builder: T, capacity: usize) -> Self {
86 let offsets_builder = Vec::with_capacity(capacity);
87 let sizes_builder = Vec::with_capacity(capacity);
88 Self {
89 offsets_builder,
90 null_buffer_builder: NullBufferBuilder::new(capacity),
91 values_builder,
92 sizes_builder,
93 field: None,
94 current_offset: OffsetSize::zero(),
95 }
96 }
97
98 pub fn with_field(self, field: impl Into<FieldRef>) -> Self {
104 Self {
105 field: Some(field.into()),
106 ..self
107 }
108 }
109}
110
111impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListViewBuilder<OffsetSize, T>
112where
113 T: 'static,
114{
115 pub fn values(&mut self) -> &mut T {
120 &mut self.values_builder
121 }
122
123 pub fn values_ref(&self) -> &T {
125 &self.values_builder
126 }
127
128 #[inline]
134 pub fn append(&mut self, is_valid: bool) {
135 self.offsets_builder.push(self.current_offset);
136 self.sizes_builder.push(
137 OffsetSize::from_usize(
138 self.values_builder.len() - self.current_offset.to_usize().unwrap(),
139 )
140 .unwrap(),
141 );
142 self.null_buffer_builder.append(is_valid);
143 self.current_offset = OffsetSize::from_usize(self.values_builder.len()).unwrap();
144 }
145
146 #[inline]
148 pub fn append_value<I, V>(&mut self, i: I)
149 where
150 T: Extend<Option<V>>,
151 I: IntoIterator<Item = Option<V>>,
152 {
153 self.extend(std::iter::once(Some(i)))
154 }
155
156 #[inline]
160 pub fn append_null(&mut self) {
161 self.offsets_builder.push(self.current_offset);
162 self.sizes_builder.push(OffsetSize::from_usize(0).unwrap());
163 self.null_buffer_builder.append_null();
164 }
165
166 #[inline]
170 pub fn append_option<I, V>(&mut self, i: Option<I>)
171 where
172 T: Extend<Option<V>>,
173 I: IntoIterator<Item = Option<V>>,
174 {
175 match i {
176 Some(i) => self.append_value(i),
177 None => self.append_null(),
178 }
179 }
180
181 pub fn finish(&mut self) -> GenericListViewArray<OffsetSize> {
183 let values = self.values_builder.finish();
184 let nulls = self.null_buffer_builder.finish();
185 let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
186 self.current_offset = OffsetSize::zero();
187
188 let offsets = ScalarBuffer::from(offsets);
190 let sizes = Buffer::from_vec(std::mem::take(&mut self.sizes_builder));
191 let sizes = ScalarBuffer::from(sizes);
192 let field = match &self.field {
193 Some(f) => f.clone(),
194 None => Arc::new(Field::new("item", values.data_type().clone(), true)),
195 };
196 GenericListViewArray::new(field, offsets, sizes, values, nulls)
197 }
198
199 pub fn finish_cloned(&self) -> GenericListViewArray<OffsetSize> {
201 let values = self.values_builder.finish_cloned();
202 let nulls = self.null_buffer_builder.finish_cloned();
203
204 let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice());
205 let offsets = ScalarBuffer::from(offsets);
207
208 let sizes = Buffer::from_slice_ref(self.sizes_builder.as_slice());
209 let sizes = ScalarBuffer::from(sizes);
210
211 let field = match &self.field {
212 Some(f) => f.clone(),
213 None => Arc::new(Field::new("item", values.data_type().clone(), true)),
214 };
215
216 GenericListViewArray::new(field, offsets, sizes, values, nulls)
217 }
218
219 pub fn offsets_slice(&self) -> &[OffsetSize] {
221 self.offsets_builder.as_slice()
222 }
223}
224
225impl<O, B, V, E> Extend<Option<V>> for GenericListViewBuilder<O, B>
226where
227 O: OffsetSizeTrait,
228 B: ArrayBuilder + Extend<E>,
229 V: IntoIterator<Item = E>,
230{
231 #[inline]
232 fn extend<T: IntoIterator<Item = Option<V>>>(&mut self, iter: T) {
233 for v in iter {
234 match v {
235 Some(elements) => {
236 self.values_builder.extend(elements);
237 self.append(true);
238 }
239 None => self.append(false),
240 }
241 }
242 }
243}
244
245#[cfg(test)]
246mod tests {
247 use super::*;
248 use crate::builder::{make_builder, Int32Builder, ListViewBuilder};
249 use crate::cast::AsArray;
250 use crate::types::Int32Type;
251 use crate::{Array, Int32Array};
252 use arrow_schema::DataType;
253
254 fn test_generic_list_view_array_builder_impl<O: OffsetSizeTrait>() {
255 let values_builder = Int32Builder::with_capacity(10);
256 let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
257
258 builder.values().append_value(0);
260 builder.values().append_value(1);
261 builder.values().append_value(2);
262 builder.append(true);
263 builder.values().append_value(3);
264 builder.values().append_value(4);
265 builder.values().append_value(5);
266 builder.append(true);
267 builder.values().append_value(6);
268 builder.values().append_value(7);
269 builder.append(true);
270 let list_array = builder.finish();
271
272 let list_values = list_array.values().as_primitive::<Int32Type>();
273 assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]);
274 assert_eq!(list_array.value_offsets(), [0, 3, 6].map(O::usize_as));
275 assert_eq!(list_array.value_sizes(), [3, 3, 2].map(O::usize_as));
276 assert_eq!(DataType::Int32, list_array.value_type());
277 assert_eq!(3, list_array.len());
278 assert_eq!(0, list_array.null_count());
279 assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]);
280 assert_eq!(O::from_usize(2).unwrap(), list_array.value_sizes()[2]);
281 for i in 0..2 {
282 assert!(list_array.is_valid(i));
283 assert!(!list_array.is_null(i));
284 }
285 }
286
287 #[test]
288 fn test_list_view_array_builder() {
289 test_generic_list_view_array_builder_impl::<i32>()
290 }
291
292 #[test]
293 fn test_large_list_view_array_builder() {
294 test_generic_list_view_array_builder_impl::<i64>()
295 }
296
297 fn test_generic_list_view_array_builder_nulls_impl<O: OffsetSizeTrait>() {
298 let values_builder = Int32Builder::with_capacity(10);
299 let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
300
301 builder.values().append_value(0);
303 builder.values().append_value(1);
304 builder.values().append_value(2);
305 builder.append(true);
306 builder.append(false);
307 builder.values().append_value(3);
308 builder.values().append_null();
309 builder.values().append_value(5);
310 builder.append(true);
311 builder.values().append_value(6);
312 builder.values().append_value(7);
313 builder.append(true);
314
315 let list_array = builder.finish();
316
317 assert_eq!(DataType::Int32, list_array.value_type());
318 assert_eq!(4, list_array.len());
319 assert_eq!(1, list_array.null_count());
320 assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]);
321 assert_eq!(O::from_usize(3).unwrap(), list_array.value_sizes()[2]);
322 }
323
324 #[test]
325 fn test_list_view_array_builder_nulls() {
326 test_generic_list_view_array_builder_nulls_impl::<i32>()
327 }
328
329 #[test]
330 fn test_large_list_view_array_builder_nulls() {
331 test_generic_list_view_array_builder_nulls_impl::<i64>()
332 }
333
334 #[test]
335 fn test_list_view_array_builder_finish() {
336 let values_builder = Int32Array::builder(5);
337 let mut builder = ListViewBuilder::new(values_builder);
338
339 builder.values().append_slice(&[1, 2, 3]);
340 builder.append(true);
341 builder.values().append_slice(&[4, 5, 6]);
342 builder.append(true);
343
344 let mut arr = builder.finish();
345 assert_eq!(2, arr.len());
346 assert!(builder.is_empty());
347
348 builder.values().append_slice(&[7, 8, 9]);
349 builder.append(true);
350 arr = builder.finish();
351 assert_eq!(1, arr.len());
352 assert!(builder.is_empty());
353 }
354
355 #[test]
356 fn test_list_view_array_builder_finish_cloned() {
357 let values_builder = Int32Array::builder(5);
358 let mut builder = ListViewBuilder::new(values_builder);
359
360 builder.values().append_slice(&[1, 2, 3]);
361 builder.append(true);
362 builder.values().append_slice(&[4, 5, 6]);
363 builder.append(true);
364
365 let mut arr = builder.finish_cloned();
366 assert_eq!(2, arr.len());
367 assert!(!builder.is_empty());
368
369 builder.values().append_slice(&[7, 8, 9]);
370 builder.append(true);
371 arr = builder.finish();
372 assert_eq!(3, arr.len());
373 assert!(builder.is_empty());
374 }
375
376 #[test]
377 fn test_list_view_list_view_array_builder() {
378 let primitive_builder = Int32Builder::with_capacity(10);
379 let values_builder = ListViewBuilder::new(primitive_builder);
380 let mut builder = ListViewBuilder::new(values_builder);
381
382 builder.values().values().append_value(1);
384 builder.values().values().append_value(2);
385 builder.values().append(true);
386 builder.values().values().append_value(3);
387 builder.values().values().append_value(4);
388 builder.values().append(true);
389 builder.append(true);
390
391 builder.values().values().append_value(5);
392 builder.values().values().append_value(6);
393 builder.values().values().append_value(7);
394 builder.values().append(true);
395 builder.values().append(false);
396 builder.values().values().append_value(8);
397 builder.values().append(true);
398 builder.append(true);
399
400 builder.append(false);
401
402 builder.values().values().append_value(9);
403 builder.values().values().append_value(10);
404 builder.values().append(true);
405 builder.append(true);
406
407 let l1 = builder.finish();
408
409 assert_eq!(4, l1.len());
410 assert_eq!(1, l1.null_count());
411
412 assert_eq!(l1.value_offsets(), &[0, 2, 5, 5]);
413 assert_eq!(l1.value_sizes(), &[2, 3, 0, 1]);
414
415 let l2 = l1.values().as_list_view::<i32>();
416
417 assert_eq!(6, l2.len());
418 assert_eq!(1, l2.null_count());
419 assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8]);
420 assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2]);
421
422 let i1 = l2.values().as_primitive::<Int32Type>();
423 assert_eq!(10, i1.len());
424 assert_eq!(0, i1.null_count());
425 assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
426 }
427
428 #[test]
429 fn test_extend() {
430 let mut builder = ListViewBuilder::new(Int32Builder::new());
431 builder.extend([
432 Some(vec![Some(1), Some(2), Some(7), None]),
433 Some(vec![]),
434 Some(vec![Some(4), Some(5)]),
435 None,
436 ]);
437
438 let array = builder.finish();
439 assert_eq!(array.value_offsets(), [0, 4, 4, 6]);
440 assert_eq!(array.value_sizes(), [4, 0, 2, 0]);
441 assert_eq!(array.null_count(), 1);
442 assert!(array.is_null(3));
443 let elements = array.values().as_primitive::<Int32Type>();
444 assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]);
445 assert_eq!(elements.null_count(), 1);
446 assert!(elements.is_null(3));
447 }
448
449 #[test]
450 fn test_boxed_primitive_array_builder() {
451 let values_builder = make_builder(&DataType::Int32, 5);
452 let mut builder = ListViewBuilder::new(values_builder);
453
454 builder
455 .values()
456 .as_any_mut()
457 .downcast_mut::<Int32Builder>()
458 .expect("should be an Int32Builder")
459 .append_slice(&[1, 2, 3]);
460 builder.append(true);
461
462 builder
463 .values()
464 .as_any_mut()
465 .downcast_mut::<Int32Builder>()
466 .expect("should be an Int32Builder")
467 .append_slice(&[4, 5, 6]);
468 builder.append(true);
469
470 let arr = builder.finish();
471 assert_eq!(2, arr.len());
472
473 let elements = arr.values().as_primitive::<Int32Type>();
474 assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
475 }
476
477 #[test]
478 fn test_boxed_list_view_list_view_array_builder() {
479 let values_builder = make_builder(
481 &DataType::ListView(Arc::new(Field::new("item", DataType::Int32, true))),
482 10,
483 );
484 test_boxed_generic_list_view_generic_list_view_array_builder::<i32>(values_builder);
485 }
486
487 #[test]
488 fn test_boxed_large_list_view_large_list_view_array_builder() {
489 let values_builder = make_builder(
491 &DataType::LargeListView(Arc::new(Field::new("item", DataType::Int32, true))),
492 10,
493 );
494 test_boxed_generic_list_view_generic_list_view_array_builder::<i64>(values_builder);
495 }
496
497 fn test_boxed_generic_list_view_generic_list_view_array_builder<O>(
498 values_builder: Box<dyn ArrayBuilder>,
499 ) where
500 O: OffsetSizeTrait + PartialEq,
501 {
502 let mut builder: GenericListViewBuilder<O, Box<dyn ArrayBuilder>> =
503 GenericListViewBuilder::<O, Box<dyn ArrayBuilder>>::new(values_builder);
504
505 builder
507 .values()
508 .as_any_mut()
509 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
510 .expect("should be an (Large)ListViewBuilder")
511 .values()
512 .as_any_mut()
513 .downcast_mut::<Int32Builder>()
514 .expect("should be an Int32Builder")
515 .append_value(1);
516 builder
517 .values()
518 .as_any_mut()
519 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
520 .expect("should be an (Large)ListViewBuilder")
521 .values()
522 .as_any_mut()
523 .downcast_mut::<Int32Builder>()
524 .expect("should be an Int32Builder")
525 .append_value(2);
526 builder
527 .values()
528 .as_any_mut()
529 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
530 .expect("should be an (Large)ListViewBuilder")
531 .append(true);
532 builder
533 .values()
534 .as_any_mut()
535 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
536 .expect("should be an (Large)ListViewBuilder")
537 .values()
538 .as_any_mut()
539 .downcast_mut::<Int32Builder>()
540 .expect("should be an Int32Builder")
541 .append_value(3);
542 builder
543 .values()
544 .as_any_mut()
545 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
546 .expect("should be an (Large)ListViewBuilder")
547 .values()
548 .as_any_mut()
549 .downcast_mut::<Int32Builder>()
550 .expect("should be an Int32Builder")
551 .append_value(4);
552 builder
553 .values()
554 .as_any_mut()
555 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
556 .expect("should be an (Large)ListViewBuilder")
557 .append(true);
558 builder.append(true);
559
560 builder
561 .values()
562 .as_any_mut()
563 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
564 .expect("should be an (Large)ListViewBuilder")
565 .values()
566 .as_any_mut()
567 .downcast_mut::<Int32Builder>()
568 .expect("should be an Int32Builder")
569 .append_value(5);
570 builder
571 .values()
572 .as_any_mut()
573 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
574 .expect("should be an (Large)ListViewBuilder")
575 .values()
576 .as_any_mut()
577 .downcast_mut::<Int32Builder>()
578 .expect("should be an Int32Builder")
579 .append_value(6);
580 builder
581 .values()
582 .as_any_mut()
583 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
584 .expect("should be an (Large)ListViewBuilder")
585 .values()
586 .as_any_mut()
587 .downcast_mut::<Int32Builder>()
588 .expect("should be an (Large)ListViewBuilder")
589 .append_value(7);
590 builder
591 .values()
592 .as_any_mut()
593 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
594 .expect("should be an (Large)ListViewBuilder")
595 .append(true);
596 builder
597 .values()
598 .as_any_mut()
599 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
600 .expect("should be an (Large)ListViewBuilder")
601 .append(false);
602 builder
603 .values()
604 .as_any_mut()
605 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
606 .expect("should be an (Large)ListViewBuilder")
607 .values()
608 .as_any_mut()
609 .downcast_mut::<Int32Builder>()
610 .expect("should be an Int32Builder")
611 .append_value(8);
612 builder
613 .values()
614 .as_any_mut()
615 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
616 .expect("should be an (Large)ListViewBuilder")
617 .append(true);
618 builder.append(true);
619
620 builder.append(false);
621
622 builder
623 .values()
624 .as_any_mut()
625 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
626 .expect("should be an (Large)ListViewBuilder")
627 .values()
628 .as_any_mut()
629 .downcast_mut::<Int32Builder>()
630 .expect("should be an Int32Builder")
631 .append_value(9);
632 builder
633 .values()
634 .as_any_mut()
635 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
636 .expect("should be an (Large)ListViewBuilder")
637 .values()
638 .as_any_mut()
639 .downcast_mut::<Int32Builder>()
640 .expect("should be an Int32Builder")
641 .append_value(10);
642 builder
643 .values()
644 .as_any_mut()
645 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
646 .expect("should be an (Large)ListViewBuilder")
647 .append(true);
648 builder.append(true);
649
650 let l1 = builder.finish();
651 assert_eq!(4, l1.len());
652 assert_eq!(1, l1.null_count());
653 assert_eq!(l1.value_offsets(), &[0, 2, 5, 5].map(O::usize_as));
654 assert_eq!(l1.value_sizes(), &[2, 3, 0, 1].map(O::usize_as));
655
656 let l2 = l1.values().as_list_view::<O>();
657 assert_eq!(6, l2.len());
658 assert_eq!(1, l2.null_count());
659 assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8].map(O::usize_as));
660 assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2].map(O::usize_as));
661
662 let i1 = l2.values().as_primitive::<Int32Type>();
663 assert_eq!(10, i1.len());
664 assert_eq!(0, i1.null_count());
665 assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
666 }
667
668 #[test]
669 fn test_with_field() {
670 let field = Arc::new(Field::new("bar", DataType::Int32, false));
671 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
672 builder.append_value([Some(1), Some(2), Some(3)]);
673 builder.append_null(); builder.append_value([Some(4)]);
675 let array = builder.finish();
676 assert_eq!(array.len(), 3);
677 assert_eq!(array.data_type(), &DataType::ListView(field.clone()));
678
679 builder.append_value([Some(4), Some(5)]);
680 let array = builder.finish();
681 assert_eq!(array.data_type(), &DataType::ListView(field));
682 assert_eq!(array.len(), 1);
683 }
684
685 #[test]
686 #[should_panic(
687 expected = r#"Non-nullable field of ListViewArray \"item\" cannot contain nulls"#
688 )]
689 fn test_checks_nullability() {
691 let field = Arc::new(Field::new("item", DataType::Int32, false));
692 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
693 builder.append_value([Some(1), None]);
694 builder.finish();
695 }
696
697 #[test]
698 #[should_panic(expected = "ListViewArray expected data type Int64 got Int32")]
699 fn test_checks_data_type() {
701 let field = Arc::new(Field::new("item", DataType::Int64, false));
702 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
703 builder.append_value([Some(1)]);
704 builder.finish();
705 }
706}