1use crate::builder::ArrayBuilder;
19use crate::{ArrayRef, GenericListViewArray, OffsetSizeTrait};
20use arrow_buffer::{Buffer, BufferBuilder, 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: BufferBuilder<OffsetSize>,
29 sizes_builder: BufferBuilder<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 = BufferBuilder::<OffsetSize>::new(capacity);
87 let sizes_builder = BufferBuilder::<OffsetSize>::new(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.append(self.current_offset);
136 self.sizes_builder.append(
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.append(self.current_offset);
162 self.sizes_builder
163 .append(OffsetSize::from_usize(0).unwrap());
164 self.null_buffer_builder.append_null();
165 }
166
167 #[inline]
171 pub fn append_option<I, V>(&mut self, i: Option<I>)
172 where
173 T: Extend<Option<V>>,
174 I: IntoIterator<Item = Option<V>>,
175 {
176 match i {
177 Some(i) => self.append_value(i),
178 None => self.append_null(),
179 }
180 }
181
182 pub fn finish(&mut self) -> GenericListViewArray<OffsetSize> {
184 let values = self.values_builder.finish();
185 let nulls = self.null_buffer_builder.finish();
186 let offsets = self.offsets_builder.finish();
187 self.current_offset = OffsetSize::zero();
188
189 let offsets = ScalarBuffer::from(offsets);
191 let sizes = self.sizes_builder.finish();
192 let sizes = ScalarBuffer::from(sizes);
193 let field = match &self.field {
194 Some(f) => f.clone(),
195 None => Arc::new(Field::new("item", values.data_type().clone(), true)),
196 };
197 GenericListViewArray::new(field, offsets, sizes, values, nulls)
198 }
199
200 pub fn finish_cloned(&self) -> GenericListViewArray<OffsetSize> {
202 let values = self.values_builder.finish_cloned();
203 let nulls = self.null_buffer_builder.finish_cloned();
204
205 let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice());
206 let offsets = ScalarBuffer::from(offsets);
208
209 let sizes = Buffer::from_slice_ref(self.sizes_builder.as_slice());
210 let sizes = ScalarBuffer::from(sizes);
211
212 let field = match &self.field {
213 Some(f) => f.clone(),
214 None => Arc::new(Field::new("item", values.data_type().clone(), true)),
215 };
216
217 GenericListViewArray::new(field, offsets, sizes, values, nulls)
218 }
219
220 pub fn offsets_slice(&self) -> &[OffsetSize] {
222 self.offsets_builder.as_slice()
223 }
224}
225
226impl<O, B, V, E> Extend<Option<V>> for GenericListViewBuilder<O, B>
227where
228 O: OffsetSizeTrait,
229 B: ArrayBuilder + Extend<E>,
230 V: IntoIterator<Item = E>,
231{
232 #[inline]
233 fn extend<T: IntoIterator<Item = Option<V>>>(&mut self, iter: T) {
234 for v in iter {
235 match v {
236 Some(elements) => {
237 self.values_builder.extend(elements);
238 self.append(true);
239 }
240 None => self.append(false),
241 }
242 }
243 }
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249 use crate::builder::{make_builder, Int32Builder, ListViewBuilder};
250 use crate::cast::AsArray;
251 use crate::types::Int32Type;
252 use crate::{Array, Int32Array};
253 use arrow_schema::DataType;
254
255 fn test_generic_list_view_array_builder_impl<O: OffsetSizeTrait>() {
256 let values_builder = Int32Builder::with_capacity(10);
257 let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
258
259 builder.values().append_value(0);
261 builder.values().append_value(1);
262 builder.values().append_value(2);
263 builder.append(true);
264 builder.values().append_value(3);
265 builder.values().append_value(4);
266 builder.values().append_value(5);
267 builder.append(true);
268 builder.values().append_value(6);
269 builder.values().append_value(7);
270 builder.append(true);
271 let list_array = builder.finish();
272
273 let list_values = list_array.values().as_primitive::<Int32Type>();
274 assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]);
275 assert_eq!(list_array.value_offsets(), [0, 3, 6].map(O::usize_as));
276 assert_eq!(list_array.value_sizes(), [3, 3, 2].map(O::usize_as));
277 assert_eq!(DataType::Int32, list_array.value_type());
278 assert_eq!(3, list_array.len());
279 assert_eq!(0, list_array.null_count());
280 assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]);
281 assert_eq!(O::from_usize(2).unwrap(), list_array.value_sizes()[2]);
282 for i in 0..2 {
283 assert!(list_array.is_valid(i));
284 assert!(!list_array.is_null(i));
285 }
286 }
287
288 #[test]
289 fn test_list_view_array_builder() {
290 test_generic_list_view_array_builder_impl::<i32>()
291 }
292
293 #[test]
294 fn test_large_list_view_array_builder() {
295 test_generic_list_view_array_builder_impl::<i64>()
296 }
297
298 fn test_generic_list_view_array_builder_nulls_impl<O: OffsetSizeTrait>() {
299 let values_builder = Int32Builder::with_capacity(10);
300 let mut builder = GenericListViewBuilder::<O, _>::new(values_builder);
301
302 builder.values().append_value(0);
304 builder.values().append_value(1);
305 builder.values().append_value(2);
306 builder.append(true);
307 builder.append(false);
308 builder.values().append_value(3);
309 builder.values().append_null();
310 builder.values().append_value(5);
311 builder.append(true);
312 builder.values().append_value(6);
313 builder.values().append_value(7);
314 builder.append(true);
315
316 let list_array = builder.finish();
317
318 assert_eq!(DataType::Int32, list_array.value_type());
319 assert_eq!(4, list_array.len());
320 assert_eq!(1, list_array.null_count());
321 assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]);
322 assert_eq!(O::from_usize(3).unwrap(), list_array.value_sizes()[2]);
323 }
324
325 #[test]
326 fn test_list_view_array_builder_nulls() {
327 test_generic_list_view_array_builder_nulls_impl::<i32>()
328 }
329
330 #[test]
331 fn test_large_list_view_array_builder_nulls() {
332 test_generic_list_view_array_builder_nulls_impl::<i64>()
333 }
334
335 #[test]
336 fn test_list_view_array_builder_finish() {
337 let values_builder = Int32Array::builder(5);
338 let mut builder = ListViewBuilder::new(values_builder);
339
340 builder.values().append_slice(&[1, 2, 3]);
341 builder.append(true);
342 builder.values().append_slice(&[4, 5, 6]);
343 builder.append(true);
344
345 let mut arr = builder.finish();
346 assert_eq!(2, arr.len());
347 assert!(builder.is_empty());
348
349 builder.values().append_slice(&[7, 8, 9]);
350 builder.append(true);
351 arr = builder.finish();
352 assert_eq!(1, arr.len());
353 assert!(builder.is_empty());
354 }
355
356 #[test]
357 fn test_list_view_array_builder_finish_cloned() {
358 let values_builder = Int32Array::builder(5);
359 let mut builder = ListViewBuilder::new(values_builder);
360
361 builder.values().append_slice(&[1, 2, 3]);
362 builder.append(true);
363 builder.values().append_slice(&[4, 5, 6]);
364 builder.append(true);
365
366 let mut arr = builder.finish_cloned();
367 assert_eq!(2, arr.len());
368 assert!(!builder.is_empty());
369
370 builder.values().append_slice(&[7, 8, 9]);
371 builder.append(true);
372 arr = builder.finish();
373 assert_eq!(3, arr.len());
374 assert!(builder.is_empty());
375 }
376
377 #[test]
378 fn test_list_view_list_view_array_builder() {
379 let primitive_builder = Int32Builder::with_capacity(10);
380 let values_builder = ListViewBuilder::new(primitive_builder);
381 let mut builder = ListViewBuilder::new(values_builder);
382
383 builder.values().values().append_value(1);
385 builder.values().values().append_value(2);
386 builder.values().append(true);
387 builder.values().values().append_value(3);
388 builder.values().values().append_value(4);
389 builder.values().append(true);
390 builder.append(true);
391
392 builder.values().values().append_value(5);
393 builder.values().values().append_value(6);
394 builder.values().values().append_value(7);
395 builder.values().append(true);
396 builder.values().append(false);
397 builder.values().values().append_value(8);
398 builder.values().append(true);
399 builder.append(true);
400
401 builder.append(false);
402
403 builder.values().values().append_value(9);
404 builder.values().values().append_value(10);
405 builder.values().append(true);
406 builder.append(true);
407
408 let l1 = builder.finish();
409
410 assert_eq!(4, l1.len());
411 assert_eq!(1, l1.null_count());
412
413 assert_eq!(l1.value_offsets(), &[0, 2, 5, 5]);
414 assert_eq!(l1.value_sizes(), &[2, 3, 0, 1]);
415
416 let l2 = l1.values().as_list_view::<i32>();
417
418 assert_eq!(6, l2.len());
419 assert_eq!(1, l2.null_count());
420 assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8]);
421 assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2]);
422
423 let i1 = l2.values().as_primitive::<Int32Type>();
424 assert_eq!(10, i1.len());
425 assert_eq!(0, i1.null_count());
426 assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
427 }
428
429 #[test]
430 fn test_extend() {
431 let mut builder = ListViewBuilder::new(Int32Builder::new());
432 builder.extend([
433 Some(vec![Some(1), Some(2), Some(7), None]),
434 Some(vec![]),
435 Some(vec![Some(4), Some(5)]),
436 None,
437 ]);
438
439 let array = builder.finish();
440 assert_eq!(array.value_offsets(), [0, 4, 4, 6]);
441 assert_eq!(array.value_sizes(), [4, 0, 2, 0]);
442 assert_eq!(array.null_count(), 1);
443 assert!(array.is_null(3));
444 let elements = array.values().as_primitive::<Int32Type>();
445 assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]);
446 assert_eq!(elements.null_count(), 1);
447 assert!(elements.is_null(3));
448 }
449
450 #[test]
451 fn test_boxed_primitive_array_builder() {
452 let values_builder = make_builder(&DataType::Int32, 5);
453 let mut builder = ListViewBuilder::new(values_builder);
454
455 builder
456 .values()
457 .as_any_mut()
458 .downcast_mut::<Int32Builder>()
459 .expect("should be an Int32Builder")
460 .append_slice(&[1, 2, 3]);
461 builder.append(true);
462
463 builder
464 .values()
465 .as_any_mut()
466 .downcast_mut::<Int32Builder>()
467 .expect("should be an Int32Builder")
468 .append_slice(&[4, 5, 6]);
469 builder.append(true);
470
471 let arr = builder.finish();
472 assert_eq!(2, arr.len());
473
474 let elements = arr.values().as_primitive::<Int32Type>();
475 assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
476 }
477
478 #[test]
479 fn test_boxed_list_view_list_view_array_builder() {
480 let values_builder = make_builder(
482 &DataType::ListView(Arc::new(Field::new("item", DataType::Int32, true))),
483 10,
484 );
485 test_boxed_generic_list_view_generic_list_view_array_builder::<i32>(values_builder);
486 }
487
488 #[test]
489 fn test_boxed_large_list_view_large_list_view_array_builder() {
490 let values_builder = make_builder(
492 &DataType::LargeListView(Arc::new(Field::new("item", DataType::Int32, true))),
493 10,
494 );
495 test_boxed_generic_list_view_generic_list_view_array_builder::<i64>(values_builder);
496 }
497
498 fn test_boxed_generic_list_view_generic_list_view_array_builder<O>(
499 values_builder: Box<dyn ArrayBuilder>,
500 ) where
501 O: OffsetSizeTrait + PartialEq,
502 {
503 let mut builder: GenericListViewBuilder<O, Box<dyn ArrayBuilder>> =
504 GenericListViewBuilder::<O, Box<dyn ArrayBuilder>>::new(values_builder);
505
506 builder
508 .values()
509 .as_any_mut()
510 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
511 .expect("should be an (Large)ListViewBuilder")
512 .values()
513 .as_any_mut()
514 .downcast_mut::<Int32Builder>()
515 .expect("should be an Int32Builder")
516 .append_value(1);
517 builder
518 .values()
519 .as_any_mut()
520 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
521 .expect("should be an (Large)ListViewBuilder")
522 .values()
523 .as_any_mut()
524 .downcast_mut::<Int32Builder>()
525 .expect("should be an Int32Builder")
526 .append_value(2);
527 builder
528 .values()
529 .as_any_mut()
530 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
531 .expect("should be an (Large)ListViewBuilder")
532 .append(true);
533 builder
534 .values()
535 .as_any_mut()
536 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
537 .expect("should be an (Large)ListViewBuilder")
538 .values()
539 .as_any_mut()
540 .downcast_mut::<Int32Builder>()
541 .expect("should be an Int32Builder")
542 .append_value(3);
543 builder
544 .values()
545 .as_any_mut()
546 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
547 .expect("should be an (Large)ListViewBuilder")
548 .values()
549 .as_any_mut()
550 .downcast_mut::<Int32Builder>()
551 .expect("should be an Int32Builder")
552 .append_value(4);
553 builder
554 .values()
555 .as_any_mut()
556 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
557 .expect("should be an (Large)ListViewBuilder")
558 .append(true);
559 builder.append(true);
560
561 builder
562 .values()
563 .as_any_mut()
564 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
565 .expect("should be an (Large)ListViewBuilder")
566 .values()
567 .as_any_mut()
568 .downcast_mut::<Int32Builder>()
569 .expect("should be an Int32Builder")
570 .append_value(5);
571 builder
572 .values()
573 .as_any_mut()
574 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
575 .expect("should be an (Large)ListViewBuilder")
576 .values()
577 .as_any_mut()
578 .downcast_mut::<Int32Builder>()
579 .expect("should be an Int32Builder")
580 .append_value(6);
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 (Large)ListViewBuilder")
590 .append_value(7);
591 builder
592 .values()
593 .as_any_mut()
594 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
595 .expect("should be an (Large)ListViewBuilder")
596 .append(true);
597 builder
598 .values()
599 .as_any_mut()
600 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
601 .expect("should be an (Large)ListViewBuilder")
602 .append(false);
603 builder
604 .values()
605 .as_any_mut()
606 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
607 .expect("should be an (Large)ListViewBuilder")
608 .values()
609 .as_any_mut()
610 .downcast_mut::<Int32Builder>()
611 .expect("should be an Int32Builder")
612 .append_value(8);
613 builder
614 .values()
615 .as_any_mut()
616 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
617 .expect("should be an (Large)ListViewBuilder")
618 .append(true);
619 builder.append(true);
620
621 builder.append(false);
622
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(9);
633 builder
634 .values()
635 .as_any_mut()
636 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
637 .expect("should be an (Large)ListViewBuilder")
638 .values()
639 .as_any_mut()
640 .downcast_mut::<Int32Builder>()
641 .expect("should be an Int32Builder")
642 .append_value(10);
643 builder
644 .values()
645 .as_any_mut()
646 .downcast_mut::<GenericListViewBuilder<O, Box<dyn ArrayBuilder>>>()
647 .expect("should be an (Large)ListViewBuilder")
648 .append(true);
649 builder.append(true);
650
651 let l1 = builder.finish();
652 assert_eq!(4, l1.len());
653 assert_eq!(1, l1.null_count());
654 assert_eq!(l1.value_offsets(), &[0, 2, 5, 5].map(O::usize_as));
655 assert_eq!(l1.value_sizes(), &[2, 3, 0, 1].map(O::usize_as));
656
657 let l2 = l1.values().as_list_view::<O>();
658 assert_eq!(6, l2.len());
659 assert_eq!(1, l2.null_count());
660 assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8].map(O::usize_as));
661 assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2].map(O::usize_as));
662
663 let i1 = l2.values().as_primitive::<Int32Type>();
664 assert_eq!(10, i1.len());
665 assert_eq!(0, i1.null_count());
666 assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
667 }
668
669 #[test]
670 fn test_with_field() {
671 let field = Arc::new(Field::new("bar", DataType::Int32, false));
672 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
673 builder.append_value([Some(1), Some(2), Some(3)]);
674 builder.append_null(); builder.append_value([Some(4)]);
676 let array = builder.finish();
677 assert_eq!(array.len(), 3);
678 assert_eq!(array.data_type(), &DataType::ListView(field.clone()));
679
680 builder.append_value([Some(4), Some(5)]);
681 let array = builder.finish();
682 assert_eq!(array.data_type(), &DataType::ListView(field));
683 assert_eq!(array.len(), 1);
684 }
685
686 #[test]
687 #[should_panic(
688 expected = r#"Non-nullable field of ListViewArray \"item\" cannot contain nulls"#
689 )]
690 fn test_checks_nullability() {
692 let field = Arc::new(Field::new("item", DataType::Int32, false));
693 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
694 builder.append_value([Some(1), None]);
695 builder.finish();
696 }
697
698 #[test]
699 #[should_panic(expected = "ListViewArray expected data type Int64 got Int32")]
700 fn test_checks_data_type() {
702 let field = Arc::new(Field::new("item", DataType::Int64, false));
703 let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone());
704 builder.append_value([Some(1)]);
705 builder.finish();
706 }
707}