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 fn finish_preserve_values(&mut self) -> ArrayRef {
76 Arc::new(self.finish_preserve_values())
77 }
78}
79
80impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListViewBuilder<OffsetSize, T> {
81 pub fn new(values_builder: T) -> Self {
83 let capacity = values_builder.len();
84 Self::with_capacity(values_builder, capacity)
85 }
86
87 pub fn with_capacity(values_builder: T, capacity: usize) -> Self {
90 let offsets_builder = Vec::with_capacity(capacity);
91 let sizes_builder = Vec::with_capacity(capacity);
92 Self {
93 offsets_builder,
94 null_buffer_builder: NullBufferBuilder::new(capacity),
95 values_builder,
96 sizes_builder,
97 field: None,
98 current_offset: OffsetSize::zero(),
99 }
100 }
101
102 pub fn with_field(self, field: impl Into<FieldRef>) -> Self {
108 Self {
109 field: Some(field.into()),
110 ..self
111 }
112 }
113}
114
115impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListViewBuilder<OffsetSize, T>
116where
117 T: 'static,
118{
119 pub fn values(&mut self) -> &mut T {
124 &mut self.values_builder
125 }
126
127 pub fn values_ref(&self) -> &T {
129 &self.values_builder
130 }
131
132 #[inline]
138 pub fn append(&mut self, is_valid: bool) {
139 self.offsets_builder.push(self.current_offset);
140 self.sizes_builder.push(
141 OffsetSize::from_usize(
142 self.values_builder.len() - self.current_offset.to_usize().unwrap(),
143 )
144 .unwrap(),
145 );
146 self.null_buffer_builder.append(is_valid);
147 self.current_offset = OffsetSize::from_usize(self.values_builder.len()).unwrap();
148 }
149
150 #[inline]
152 pub fn append_value<I, V>(&mut self, i: I)
153 where
154 T: Extend<Option<V>>,
155 I: IntoIterator<Item = Option<V>>,
156 {
157 self.extend(std::iter::once(Some(i)))
158 }
159
160 #[inline]
164 pub fn append_null(&mut self) {
165 self.offsets_builder.push(self.current_offset);
166 self.sizes_builder.push(OffsetSize::from_usize(0).unwrap());
167 self.null_buffer_builder.append_null();
168 }
169
170 #[inline]
174 pub fn append_option<I, V>(&mut self, i: Option<I>)
175 where
176 T: Extend<Option<V>>,
177 I: IntoIterator<Item = Option<V>>,
178 {
179 match i {
180 Some(i) => self.append_value(i),
181 None => self.append_null(),
182 }
183 }
184
185 pub fn finish(&mut self) -> GenericListViewArray<OffsetSize> {
187 let values = self.values_builder.finish();
188 let nulls = self.null_buffer_builder.finish();
189 let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
190 self.current_offset = OffsetSize::zero();
191
192 let offsets = ScalarBuffer::from(offsets);
194 let sizes = Buffer::from_vec(std::mem::take(&mut self.sizes_builder));
195 let sizes = ScalarBuffer::from(sizes);
196 let field = match &self.field {
197 Some(f) => f.clone(),
198 None => Arc::new(Field::new("item", values.data_type().clone(), true)),
199 };
200 GenericListViewArray::new(field, offsets, sizes, values, nulls)
201 }
202
203 pub fn finish_cloned(&self) -> GenericListViewArray<OffsetSize> {
205 let values = self.values_builder.finish_cloned();
206 let nulls = self.null_buffer_builder.finish_cloned();
207
208 let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice());
209 let offsets = ScalarBuffer::from(offsets);
211
212 let sizes = Buffer::from_slice_ref(self.sizes_builder.as_slice());
213 let sizes = ScalarBuffer::from(sizes);
214
215 let field = match &self.field {
216 Some(f) => f.clone(),
217 None => Arc::new(Field::new("item", values.data_type().clone(), true)),
218 };
219
220 GenericListViewArray::new(field, offsets, sizes, values, nulls)
221 }
222
223 fn finish_preserve_values(&mut self) -> GenericListViewArray<OffsetSize> {
224 let values = self.values_builder.finish_preserve_values();
225 let nulls = self.null_buffer_builder.finish();
226 let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
227 self.current_offset = OffsetSize::zero();
228
229 let offsets = ScalarBuffer::from(offsets);
231 let sizes = Buffer::from_vec(std::mem::take(&mut self.sizes_builder));
232 let sizes = ScalarBuffer::from(sizes);
233 let field = match &self.field {
234 Some(f) => f.clone(),
235 None => Arc::new(Field::new("item", values.data_type().clone(), true)),
236 };
237 GenericListViewArray::new(field, offsets, sizes, values, nulls)
238 }
239
240 pub fn offsets_slice(&self) -> &[OffsetSize] {
242 self.offsets_builder.as_slice()
243 }
244}
245
246impl<O, B, V, E> Extend<Option<V>> for GenericListViewBuilder<O, B>
247where
248 O: OffsetSizeTrait,
249 B: ArrayBuilder + Extend<E>,
250 V: IntoIterator<Item = E>,
251{
252 #[inline]
253 fn extend<T: IntoIterator<Item = Option<V>>>(&mut self, iter: T) {
254 for v in iter {
255 match v {
256 Some(elements) => {
257 self.values_builder.extend(elements);
258 self.append(true);
259 }
260 None => self.append(false),
261 }
262 }
263 }
264}
265
266#[cfg(test)]
267mod tests {
268 use super::*;
269 use crate::builder::{Int32Builder, ListViewBuilder, make_builder, tests::PreserveValuesMock};
270 use crate::cast::AsArray;
271 use crate::types::Int32Type;
272 use crate::{Array, Int32Array};
273 use arrow_schema::DataType;
274
275 fn test_generic_list_view_array_builder_impl<O: OffsetSizeTrait>() {
276 let values_builder = Int32Builder::with_capacity(10);
277 let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
278
279 builder.values().append_value(0);
281 builder.values().append_value(1);
282 builder.values().append_value(2);
283 builder.append(true);
284 builder.values().append_value(3);
285 builder.values().append_value(4);
286 builder.values().append_value(5);
287 builder.append(true);
288 builder.values().append_value(6);
289 builder.values().append_value(7);
290 builder.append(true);
291 let list_array = builder.finish();
292
293 let list_values = list_array.values().as_primitive::<Int32Type>();
294 assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]);
295 assert_eq!(list_array.value_offsets(), [0, 3, 6].map(O::usize_as));
296 assert_eq!(list_array.value_sizes(), [3, 3, 2].map(O::usize_as));
297 assert_eq!(DataType::Int32, list_array.value_type());
298 assert_eq!(3, list_array.len());
299 assert_eq!(0, list_array.null_count());
300 assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]);
301 assert_eq!(O::from_usize(2).unwrap(), list_array.value_sizes()[2]);
302 for i in 0..2 {
303 assert!(list_array.is_valid(i));
304 assert!(!list_array.is_null(i));
305 }
306 }
307
308 #[test]
309 fn test_list_view_array_builder() {
310 test_generic_list_view_array_builder_impl::<i32>()
311 }
312
313 #[test]
314 fn test_large_list_view_array_builder() {
315 test_generic_list_view_array_builder_impl::<i64>()
316 }
317
318 fn test_generic_list_view_array_builder_nulls_impl<O: OffsetSizeTrait>() {
319 let values_builder = Int32Builder::with_capacity(10);
320 let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
321
322 builder.values().append_value(0);
324 builder.values().append_value(1);
325 builder.values().append_value(2);
326 builder.append(true);
327 builder.append(false);
328 builder.values().append_value(3);
329 builder.values().append_null();
330 builder.values().append_value(5);
331 builder.append(true);
332 builder.values().append_value(6);
333 builder.values().append_value(7);
334 builder.append(true);
335
336 let list_array = builder.finish();
337
338 assert_eq!(DataType::Int32, list_array.value_type());
339 assert_eq!(4, list_array.len());
340 assert_eq!(1, list_array.null_count());
341 assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]);
342 assert_eq!(O::from_usize(3).unwrap(), list_array.value_sizes()[2]);
343 }
344
345 #[test]
346 fn test_list_view_array_builder_nulls() {
347 test_generic_list_view_array_builder_nulls_impl::<i32>()
348 }
349
350 #[test]
351 fn test_large_list_view_array_builder_nulls() {
352 test_generic_list_view_array_builder_nulls_impl::<i64>()
353 }
354
355 #[test]
356 fn test_list_view_array_builder_finish() {
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();
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!(1, arr.len());
373 assert!(builder.is_empty());
374 }
375
376 #[test]
377 fn test_list_view_array_builder_finish_cloned() {
378 let values_builder = Int32Array::builder(5);
379 let mut builder = ListViewBuilder::new(values_builder);
380
381 builder.values().append_slice(&[1, 2, 3]);
382 builder.append(true);
383 builder.values().append_slice(&[4, 5, 6]);
384 builder.append(true);
385
386 let mut arr = builder.finish_cloned();
387 assert_eq!(2, arr.len());
388 assert!(!builder.is_empty());
389
390 builder.values().append_slice(&[7, 8, 9]);
391 builder.append(true);
392 arr = builder.finish();
393 assert_eq!(3, arr.len());
394 assert!(builder.is_empty());
395 }
396
397 #[test]
398 fn test_list_view_list_view_array_builder() {
399 let primitive_builder = Int32Builder::with_capacity(10);
400 let values_builder = ListViewBuilder::new(primitive_builder);
401 let mut builder = ListViewBuilder::new(values_builder);
402
403 builder.values().values().append_value(1);
405 builder.values().values().append_value(2);
406 builder.values().append(true);
407 builder.values().values().append_value(3);
408 builder.values().values().append_value(4);
409 builder.values().append(true);
410 builder.append(true);
411
412 builder.values().values().append_value(5);
413 builder.values().values().append_value(6);
414 builder.values().values().append_value(7);
415 builder.values().append(true);
416 builder.values().append(false);
417 builder.values().values().append_value(8);
418 builder.values().append(true);
419 builder.append(true);
420
421 builder.append(false);
422
423 builder.values().values().append_value(9);
424 builder.values().values().append_value(10);
425 builder.values().append(true);
426 builder.append(true);
427
428 let l1 = builder.finish();
429
430 assert_eq!(4, l1.len());
431 assert_eq!(1, l1.null_count());
432
433 assert_eq!(l1.value_offsets(), &[0, 2, 5, 5]);
434 assert_eq!(l1.value_sizes(), &[2, 3, 0, 1]);
435
436 let l2 = l1.values().as_list_view::<i32>();
437
438 assert_eq!(6, l2.len());
439 assert_eq!(1, l2.null_count());
440 assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8]);
441 assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2]);
442
443 let i1 = l2.values().as_primitive::<Int32Type>();
444 assert_eq!(10, i1.len());
445 assert_eq!(0, i1.null_count());
446 assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
447 }
448
449 #[test]
450 fn test_extend() {
451 let mut builder = ListViewBuilder::new(Int32Builder::new());
452 builder.extend([
453 Some(vec![Some(1), Some(2), Some(7), None]),
454 Some(vec![]),
455 Some(vec![Some(4), Some(5)]),
456 None,
457 ]);
458
459 let array = builder.finish();
460 assert_eq!(array.value_offsets(), [0, 4, 4, 6]);
461 assert_eq!(array.value_sizes(), [4, 0, 2, 0]);
462 assert_eq!(array.null_count(), 1);
463 assert!(array.is_null(3));
464 let elements = array.values().as_primitive::<Int32Type>();
465 assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]);
466 assert_eq!(elements.null_count(), 1);
467 assert!(elements.is_null(3));
468 }
469
470 #[test]
471 fn test_boxed_primitive_array_builder() {
472 let values_builder = make_builder(&DataType::Int32, 5);
473 let mut builder = ListViewBuilder::new(values_builder);
474
475 builder
476 .values()
477 .as_any_mut()
478 .downcast_mut::<Int32Builder>()
479 .expect("should be an Int32Builder")
480 .append_slice(&[1, 2, 3]);
481 builder.append(true);
482
483 builder
484 .values()
485 .as_any_mut()
486 .downcast_mut::<Int32Builder>()
487 .expect("should be an Int32Builder")
488 .append_slice(&[4, 5, 6]);
489 builder.append(true);
490
491 let arr = builder.finish();
492 assert_eq!(2, arr.len());
493
494 let elements = arr.values().as_primitive::<Int32Type>();
495 assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
496 }
497
498 #[test]
499 fn test_boxed_list_view_list_view_array_builder() {
500 let values_builder = make_builder(
502 &DataType::ListView(Arc::new(Field::new("item", DataType::Int32, true))),
503 10,
504 );
505 test_boxed_generic_list_view_generic_list_view_array_builder::<i32>(values_builder);
506 }
507
508 #[test]
509 fn test_boxed_large_list_view_large_list_view_array_builder() {
510 let values_builder = make_builder(
512 &DataType::LargeListView(Arc::new(Field::new("item", DataType::Int32, true))),
513 10,
514 );
515 test_boxed_generic_list_view_generic_list_view_array_builder::<i64>(values_builder);
516 }
517
518 fn test_boxed_generic_list_view_generic_list_view_array_builder<O>(
519 values_builder: Box<dyn ArrayBuilder>,
520 ) where
521 O: OffsetSizeTrait + PartialEq,
522 {
523 let mut builder: GenericListViewBuilder<O, Box<dyn ArrayBuilder>> =
524 GenericListViewBuilder::<O, Box<dyn ArrayBuilder>>::new(values_builder);
525
526 builder
528 .values()
529 .as_any_mut()
530 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
531 .expect("should be an (Large)ListViewBuilder")
532 .values()
533 .as_any_mut()
534 .downcast_mut::<Int32Builder>()
535 .expect("should be an Int32Builder")
536 .append_value(1);
537 builder
538 .values()
539 .as_any_mut()
540 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
541 .expect("should be an (Large)ListViewBuilder")
542 .values()
543 .as_any_mut()
544 .downcast_mut::<Int32Builder>()
545 .expect("should be an Int32Builder")
546 .append_value(2);
547 builder
548 .values()
549 .as_any_mut()
550 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
551 .expect("should be an (Large)ListViewBuilder")
552 .append(true);
553 builder
554 .values()
555 .as_any_mut()
556 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
557 .expect("should be an (Large)ListViewBuilder")
558 .values()
559 .as_any_mut()
560 .downcast_mut::<Int32Builder>()
561 .expect("should be an Int32Builder")
562 .append_value(3);
563 builder
564 .values()
565 .as_any_mut()
566 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
567 .expect("should be an (Large)ListViewBuilder")
568 .values()
569 .as_any_mut()
570 .downcast_mut::<Int32Builder>()
571 .expect("should be an Int32Builder")
572 .append_value(4);
573 builder
574 .values()
575 .as_any_mut()
576 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
577 .expect("should be an (Large)ListViewBuilder")
578 .append(true);
579 builder.append(true);
580
581 builder
582 .values()
583 .as_any_mut()
584 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
585 .expect("should be an (Large)ListViewBuilder")
586 .values()
587 .as_any_mut()
588 .downcast_mut::<Int32Builder>()
589 .expect("should be an Int32Builder")
590 .append_value(5);
591 builder
592 .values()
593 .as_any_mut()
594 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
595 .expect("should be an (Large)ListViewBuilder")
596 .values()
597 .as_any_mut()
598 .downcast_mut::<Int32Builder>()
599 .expect("should be an Int32Builder")
600 .append_value(6);
601 builder
602 .values()
603 .as_any_mut()
604 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
605 .expect("should be an (Large)ListViewBuilder")
606 .values()
607 .as_any_mut()
608 .downcast_mut::<Int32Builder>()
609 .expect("should be an (Large)ListViewBuilder")
610 .append_value(7);
611 builder
612 .values()
613 .as_any_mut()
614 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
615 .expect("should be an (Large)ListViewBuilder")
616 .append(true);
617 builder
618 .values()
619 .as_any_mut()
620 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
621 .expect("should be an (Large)ListViewBuilder")
622 .append(false);
623 builder
624 .values()
625 .as_any_mut()
626 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
627 .expect("should be an (Large)ListViewBuilder")
628 .values()
629 .as_any_mut()
630 .downcast_mut::<Int32Builder>()
631 .expect("should be an Int32Builder")
632 .append_value(8);
633 builder
634 .values()
635 .as_any_mut()
636 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
637 .expect("should be an (Large)ListViewBuilder")
638 .append(true);
639 builder.append(true);
640
641 builder.append(false);
642
643 builder
644 .values()
645 .as_any_mut()
646 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
647 .expect("should be an (Large)ListViewBuilder")
648 .values()
649 .as_any_mut()
650 .downcast_mut::<Int32Builder>()
651 .expect("should be an Int32Builder")
652 .append_value(9);
653 builder
654 .values()
655 .as_any_mut()
656 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
657 .expect("should be an (Large)ListViewBuilder")
658 .values()
659 .as_any_mut()
660 .downcast_mut::<Int32Builder>()
661 .expect("should be an Int32Builder")
662 .append_value(10);
663 builder
664 .values()
665 .as_any_mut()
666 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
667 .expect("should be an (Large)ListViewBuilder")
668 .append(true);
669 builder.append(true);
670
671 let l1 = builder.finish();
672 assert_eq!(4, l1.len());
673 assert_eq!(1, l1.null_count());
674 assert_eq!(l1.value_offsets(), &[0, 2, 5, 5].map(O::usize_as));
675 assert_eq!(l1.value_sizes(), &[2, 3, 0, 1].map(O::usize_as));
676
677 let l2 = l1.values().as_list_view::<O>();
678 assert_eq!(6, l2.len());
679 assert_eq!(1, l2.null_count());
680 assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8].map(O::usize_as));
681 assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2].map(O::usize_as));
682
683 let i1 = l2.values().as_primitive::<Int32Type>();
684 assert_eq!(10, i1.len());
685 assert_eq!(0, i1.null_count());
686 assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
687 }
688
689 #[test]
690 fn test_with_field() {
691 let field = Arc::new(Field::new("bar", DataType::Int32, false));
692 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
693 builder.append_value([Some(1), Some(2), Some(3)]);
694 builder.append_null(); builder.append_value([Some(4)]);
696 let array = builder.finish();
697 assert_eq!(array.len(), 3);
698 assert_eq!(array.data_type(), &DataType::ListView(field.clone()));
699
700 builder.append_value([Some(4), Some(5)]);
701 let array = builder.finish();
702 assert_eq!(array.data_type(), &DataType::ListView(field));
703 assert_eq!(array.len(), 1);
704 }
705
706 #[test]
707 #[should_panic(
708 expected = r#"Non-nullable field of ListViewArray \"item\" cannot contain nulls"#
709 )]
710 fn test_checks_nullability() {
712 let field = Arc::new(Field::new("item", DataType::Int32, false));
713 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
714 builder.append_value([Some(1), None]);
715 builder.finish();
716 }
717
718 #[test]
719 #[should_panic(expected = "ListViewArray expected data type Int64 got Int32")]
720 fn test_checks_data_type() {
722 let field = Arc::new(Field::new("item", DataType::Int64, false));
723 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
724 builder.append_value([Some(1)]);
725 builder.finish();
726 }
727
728 #[test]
729 fn test_finish_preserve_values() {
730 let mut builder = ListViewBuilder::new(PreserveValuesMock::default());
731
732 builder.values().inner.append_value(1);
733 builder.append(true);
734
735 let arr = builder.finish_preserve_values();
736
737 assert_eq!(1, arr.len());
738 assert_eq!(1, builder.values().called);
739 }
740}