arrow_array/builder/
generic_list_builder.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::builder::ArrayBuilder;
6use crate::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
7use arrow_buffer::NullBufferBuilder;
8use arrow_buffer::{Buffer, OffsetBuffer};
9use arrow_schema::{Field, FieldRef};
10use std::any::Any;
11use std::sync::Arc;
12
13/// Builder for [`GenericListArray`]
14///
15/// Use [`ListBuilder`] to build [`ListArray`]s and [`LargeListBuilder`] to build [`LargeListArray`]s.
16///
17/// # Example
18///
19/// Here is code that constructs a ListArray with the contents:
20/// `[[A,B,C], [], NULL, [D], [NULL, F]]`
21///
22/// ```
23/// # use std::sync::Arc;
24/// # use arrow_array::{builder::ListBuilder, builder::StringBuilder, ArrayRef, StringArray, Array};
25/// #
26/// let values_builder = StringBuilder::new();
27/// let mut builder = ListBuilder::new(values_builder);
28///
29/// // [A, B, C]
30/// builder.values().append_value("A");
31/// builder.values().append_value("B");
32/// builder.values().append_value("C");
33/// builder.append(true);
34///
35/// // [ ] (empty list)
36/// builder.append(true);
37///
38/// // Null
39/// builder.append(false);
40///
41/// // [D]
42/// builder.values().append_value("D");
43/// builder.append(true);
44///
45/// // [NULL, F]
46/// builder.values().append_null();
47/// builder.values().append_value("F");
48/// builder.append(true);
49///
50/// // Build the array
51/// let array = builder.finish();
52///
53/// // Values is a string array
54/// // "A", "B" "C", "?", "D", NULL, "F"
55/// assert_eq!(
56///   array.values().as_ref(),
57///   &StringArray::from(vec![
58///     Some("A"), Some("B"), Some("C"),
59///     Some("D"), None, Some("F")
60///   ])
61/// );
62///
63/// // Offsets are indexes into the values array
64/// assert_eq!(
65///   array.value_offsets(),
66///   &[0, 3, 3, 3, 4, 6]
67/// );
68/// ```
69///
70/// [`ListBuilder`]: crate::builder::ListBuilder
71/// [`ListArray`]: crate::array::ListArray
72/// [`LargeListBuilder`]: crate::builder::LargeListBuilder
73/// [`LargeListArray`]: crate::array::LargeListArray
74#[derive(Debug)]
75pub struct GenericListBuilder<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> {
76    offsets_builder: Vec<OffsetSize>,
77    null_buffer_builder: NullBufferBuilder,
78    values_builder: T,
79    field: Option<FieldRef>,
80}
81
82impl<O: OffsetSizeTrait, T: ArrayBuilder + Default> Default for GenericListBuilder<O, T> {
83    fn default() -> Self {
84        Self::new(T::default())
85    }
86}
87
88impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListBuilder<OffsetSize, T> {
89    /// Creates a new [`GenericListBuilder`] from a given values array builder
90    pub fn new(values_builder: T) -> Self {
91        let capacity = values_builder.len();
92        Self::with_capacity(values_builder, capacity)
93    }
94
95    /// Creates a new [`GenericListBuilder`] from a given values array builder
96    /// `capacity` is the number of items to pre-allocate space for in this builder
97    pub fn with_capacity(values_builder: T, capacity: usize) -> Self {
98        let mut offsets_builder = Vec::with_capacity(capacity + 1);
99        offsets_builder.push(OffsetSize::zero());
100        Self {
101            offsets_builder,
102            null_buffer_builder: NullBufferBuilder::new(capacity),
103            values_builder,
104            field: None,
105        }
106    }
107
108    /// Override the field passed to [`GenericListArray::new`]
109    ///
110    /// By default a nullable field is created with the name `item`
111    ///
112    /// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
113    /// field's data type does not match that of `T`
114    pub fn with_field(self, field: impl Into<FieldRef>) -> Self {
115        Self {
116            field: Some(field.into()),
117            ..self
118        }
119    }
120}
121
122impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> ArrayBuilder
123    for GenericListBuilder<OffsetSize, T>
124where
125    T: 'static,
126{
127    /// Returns the builder as a non-mutable `Any` reference.
128    fn as_any(&self) -> &dyn Any {
129        self
130    }
131
132    /// Returns the builder as a mutable `Any` reference.
133    fn as_any_mut(&mut self) -> &mut dyn Any {
134        self
135    }
136
137    /// Returns the boxed builder as a box of `Any`.
138    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
139        self
140    }
141
142    /// Returns the number of array slots in the builder
143    fn len(&self) -> usize {
144        self.null_buffer_builder.len()
145    }
146
147    /// Builds the array and reset this builder.
148    fn finish(&mut self) -> ArrayRef {
149        Arc::new(self.finish())
150    }
151
152    /// Builds the array without resetting the builder.
153    fn finish_cloned(&self) -> ArrayRef {
154        Arc::new(self.finish_cloned())
155    }
156
157    fn finish_preserve_values(&mut self) -> ArrayRef {
158        Arc::new(self.finish_preserve_values())
159    }
160}
161
162impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListBuilder<OffsetSize, T>
163where
164    T: 'static,
165{
166    /// Returns the child array builder as a mutable reference.
167    ///
168    /// This mutable reference can be used to append values into the child array builder,
169    /// but you must call [`append`](#method.append) to delimit each distinct list value.
170    pub fn values(&mut self) -> &mut T {
171        &mut self.values_builder
172    }
173
174    /// Returns the child array builder as an immutable reference
175    pub fn values_ref(&self) -> &T {
176        &self.values_builder
177    }
178
179    /// Finish the current variable-length list array slot
180    ///
181    /// # Panics
182    ///
183    /// Panics if the length of [`Self::values`] exceeds `OffsetSize::MAX`
184    #[inline]
185    pub fn append(&mut self, is_valid: bool) {
186        self.offsets_builder.push(self.next_offset());
187        self.null_buffer_builder.append(is_valid);
188    }
189
190    /// Returns the next offset
191    ///
192    /// # Panics
193    ///
194    /// Panics if the length of [`Self::values`] exceeds `OffsetSize::MAX`
195    #[inline]
196    fn next_offset(&self) -> OffsetSize {
197        OffsetSize::from_usize(self.values_builder.len()).unwrap()
198    }
199
200    /// Append a value to this [`GenericListBuilder`]
201    ///
202    /// ```
203    /// # use arrow_array::builder::{Int32Builder, ListBuilder};
204    /// # use arrow_array::cast::AsArray;
205    /// # use arrow_array::{Array, Int32Array};
206    /// # use arrow_array::types::Int32Type;
207    /// let mut builder = ListBuilder::new(Int32Builder::new());
208    ///
209    /// builder.append_value([Some(1), Some(2), Some(3)]);
210    /// builder.append_value([]);
211    /// builder.append_value([None]);
212    ///
213    /// let array = builder.finish();
214    /// assert_eq!(array.len(), 3);
215    ///
216    /// assert_eq!(array.value_offsets(), &[0, 3, 3, 4]);
217    /// let values = array.values().as_primitive::<Int32Type>();
218    /// assert_eq!(values, &Int32Array::from(vec![Some(1), Some(2), Some(3), None]));
219    /// ```
220    ///
221    /// This is an alternative API to appending directly to [`Self::values`] and
222    /// delimiting the result with [`Self::append`]
223    ///
224    /// ```
225    /// # use arrow_array::builder::{Int32Builder, ListBuilder};
226    /// # use arrow_array::cast::AsArray;
227    /// # use arrow_array::{Array, Int32Array};
228    /// # use arrow_array::types::Int32Type;
229    /// let mut builder = ListBuilder::new(Int32Builder::new());
230    ///
231    /// builder.values().append_value(1);
232    /// builder.values().append_value(2);
233    /// builder.values().append_value(3);
234    /// builder.append(true);
235    /// builder.append(true);
236    /// builder.values().append_null();
237    /// builder.append(true);
238    ///
239    /// let array = builder.finish();
240    /// assert_eq!(array.len(), 3);
241    ///
242    /// assert_eq!(array.value_offsets(), &[0, 3, 3, 4]);
243    /// let values = array.values().as_primitive::<Int32Type>();
244    /// assert_eq!(values, &Int32Array::from(vec![Some(1), Some(2), Some(3), None]));
245    /// ```
246    #[inline]
247    pub fn append_value<I, V>(&mut self, i: I)
248    where
249        T: Extend<Option<V>>,
250        I: IntoIterator<Item = Option<V>>,
251    {
252        self.extend(std::iter::once(Some(i)))
253    }
254
255    /// Append a null to this [`GenericListBuilder`]
256    ///
257    /// See [`Self::append_value`] for an example use.
258    #[inline]
259    pub fn append_null(&mut self) {
260        self.offsets_builder.push(self.next_offset());
261        self.null_buffer_builder.append_null();
262    }
263
264    /// Appends `n` `null`s into the builder.
265    #[inline]
266    pub fn append_nulls(&mut self, n: usize) {
267        let next_offset = self.next_offset();
268        self.offsets_builder
269            .extend(std::iter::repeat_n(next_offset, n));
270        self.null_buffer_builder.append_n_nulls(n);
271    }
272
273    /// Appends an optional value into this [`GenericListBuilder`]
274    ///
275    /// If `Some` calls [`Self::append_value`] otherwise calls [`Self::append_null`]
276    #[inline]
277    pub fn append_option<I, V>(&mut self, i: Option<I>)
278    where
279        T: Extend<Option<V>>,
280        I: IntoIterator<Item = Option<V>>,
281    {
282        match i {
283            Some(i) => self.append_value(i),
284            None => self.append_null(),
285        }
286    }
287
288    /// Builds the [`GenericListArray`] and reset this builder.
289    ///
290    /// # Panics
291    ///
292    /// Panics if the field set with [`Self::with_field`] does not match the data type
293    /// of the values builder
294    pub fn finish(&mut self) -> GenericListArray<OffsetSize> {
295        let values = self.values_builder.finish();
296        let nulls = self.null_buffer_builder.finish();
297
298        let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
299        // Safety: Safe by construction
300        let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
301        self.offsets_builder.push(OffsetSize::zero());
302
303        let field = match &self.field {
304            Some(f) => f.clone(),
305            None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
306        };
307
308        GenericListArray::new(field, offsets, values, nulls)
309    }
310
311    /// Builds the [`GenericListArray`] without resetting the builder.
312    ///
313    /// # Panics
314    ///
315    /// Panics if the field set with [`Self::with_field`] does not match the data type
316    /// of the values builder
317    pub fn finish_cloned(&self) -> GenericListArray<OffsetSize> {
318        let values = self.values_builder.finish_cloned();
319        let nulls = self.null_buffer_builder.finish_cloned();
320
321        let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice());
322        // Safety: safe by construction
323        let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
324
325        let field = match &self.field {
326            Some(f) => f.clone(),
327            None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
328        };
329
330        GenericListArray::new(field, offsets, values, nulls)
331    }
332
333    fn finish_preserve_values(&mut self) -> GenericListArray<OffsetSize> {
334        let values = self.values_builder.finish_preserve_values();
335        let nulls = self.null_buffer_builder.finish();
336
337        let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
338        // Safety: Safe by construction
339        let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
340        self.offsets_builder.push(OffsetSize::zero());
341
342        let field = match &self.field {
343            Some(f) => f.clone(),
344            None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
345        };
346
347        GenericListArray::new(field, offsets, values, nulls)
348    }
349
350    /// Returns the current offsets buffer as a slice
351    pub fn offsets_slice(&self) -> &[OffsetSize] {
352        self.offsets_builder.as_slice()
353    }
354
355    /// Returns the current offsets buffer capacity, in offsets.
356    pub fn offsets_capacity(&self) -> usize {
357        self.offsets_builder.capacity()
358    }
359
360    /// Returns the current null buffer as a slice
361    pub fn validity_slice(&self) -> Option<&[u8]> {
362        self.null_buffer_builder.as_slice()
363    }
364
365    /// Returns the current null buffer allocated capacity, in bytes.
366    pub fn validity_capacity(&self) -> usize {
367        self.null_buffer_builder.allocated_size()
368    }
369}
370
371impl<O, B, V, E> Extend<Option<V>> for GenericListBuilder<O, B>
372where
373    O: OffsetSizeTrait,
374    B: ArrayBuilder + Extend<E>,
375    V: IntoIterator<Item = E>,
376{
377    #[inline]
378    fn extend<T: IntoIterator<Item = Option<V>>>(&mut self, iter: T) {
379        for v in iter {
380            match v {
381                Some(elements) => {
382                    self.values_builder.extend(elements);
383                    self.append(true);
384                }
385                None => self.append(false),
386            }
387        }
388    }
389}
390
391#[cfg(test)]
392mod tests {
393    use super::*;
394    use crate::Int32Array;
395    use crate::builder::{Int32Builder, ListBuilder, make_builder, tests::PreserveValuesMock};
396    use crate::cast::AsArray;
397    use crate::types::Int32Type;
398    use arrow_schema::DataType;
399
400    fn _test_generic_list_array_builder<O: OffsetSizeTrait>() {
401        let values_builder = Int32Builder::with_capacity(10);
402        let mut builder = GenericListBuilder::<O, _>::new(values_builder);
403
404        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
405        builder.values().append_value(0);
406        builder.values().append_value(1);
407        builder.values().append_value(2);
408        builder.append(true);
409        builder.values().append_value(3);
410        builder.values().append_value(4);
411        builder.values().append_value(5);
412        builder.append(true);
413        builder.values().append_value(6);
414        builder.values().append_value(7);
415        builder.append(true);
416        let list_array = builder.finish();
417
418        let list_values = list_array.values().as_primitive::<Int32Type>();
419        assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]);
420        assert_eq!(list_array.value_offsets(), [0, 3, 6, 8].map(O::usize_as));
421        assert_eq!(DataType::Int32, list_array.value_type());
422        assert_eq!(3, list_array.len());
423        assert_eq!(0, list_array.null_count());
424        assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]);
425        assert_eq!(O::from_usize(2).unwrap(), list_array.value_length(2));
426        for i in 0..3 {
427            assert!(list_array.is_valid(i));
428            assert!(!list_array.is_null(i));
429        }
430    }
431
432    #[test]
433    fn test_list_array_builder() {
434        _test_generic_list_array_builder::<i32>()
435    }
436
437    #[test]
438    fn test_large_list_array_builder() {
439        _test_generic_list_array_builder::<i64>()
440    }
441
442    fn _test_generic_list_array_builder_nulls<O: OffsetSizeTrait>() {
443        let values_builder = Int32Builder::with_capacity(10);
444        let mut builder = GenericListBuilder::<O, _>::new(values_builder);
445
446        //  [[0, 1, 2], null, [3, null, 5], [6, 7], null, null, [8]]
447        builder.values().append_value(0);
448        builder.values().append_value(1);
449        builder.values().append_value(2);
450        builder.append(true);
451        builder.append(false);
452        builder.values().append_value(3);
453        builder.values().append_null();
454        builder.values().append_value(5);
455        builder.append(true);
456        builder.values().append_value(6);
457        builder.values().append_value(7);
458        builder.append(true);
459        builder.append_nulls(2);
460        builder.values().append_value(8);
461        builder.append(true);
462
463        let list_array = builder.finish();
464
465        assert_eq!(DataType::Int32, list_array.value_type());
466        assert_eq!(7, list_array.len());
467        assert_eq!(3, list_array.null_count());
468        assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]);
469        assert_eq!(O::from_usize(9).unwrap(), list_array.value_offsets()[7]);
470        assert_eq!(O::from_usize(3).unwrap(), list_array.value_length(2));
471        assert!(list_array.is_null(4));
472        assert!(list_array.is_null(5));
473    }
474
475    #[test]
476    fn test_list_array_builder_nulls() {
477        _test_generic_list_array_builder_nulls::<i32>()
478    }
479
480    #[test]
481    fn test_large_list_array_builder_nulls() {
482        _test_generic_list_array_builder_nulls::<i64>()
483    }
484
485    #[test]
486    fn test_list_array_builder_finish() {
487        let values_builder = Int32Array::builder(5);
488        let mut builder = ListBuilder::new(values_builder);
489
490        builder.values().append_slice(&[1, 2, 3]);
491        builder.append(true);
492        builder.values().append_slice(&[4, 5, 6]);
493        builder.append(true);
494
495        let mut arr = builder.finish();
496        assert_eq!(2, arr.len());
497        assert!(builder.is_empty());
498
499        builder.values().append_slice(&[7, 8, 9]);
500        builder.append(true);
501        arr = builder.finish();
502        assert_eq!(1, arr.len());
503        assert!(builder.is_empty());
504    }
505
506    #[test]
507    fn test_list_array_builder_finish_cloned() {
508        let values_builder = Int32Array::builder(5);
509        let mut builder = ListBuilder::new(values_builder);
510
511        builder.values().append_slice(&[1, 2, 3]);
512        builder.append(true);
513        builder.values().append_slice(&[4, 5, 6]);
514        builder.append(true);
515
516        let mut arr = builder.finish_cloned();
517        assert_eq!(2, arr.len());
518        assert!(!builder.is_empty());
519
520        builder.values().append_slice(&[7, 8, 9]);
521        builder.append(true);
522        arr = builder.finish();
523        assert_eq!(3, arr.len());
524        assert!(builder.is_empty());
525    }
526
527    #[test]
528    fn test_list_list_array_builder() {
529        let primitive_builder = Int32Builder::with_capacity(10);
530        let values_builder = ListBuilder::new(primitive_builder);
531        let mut builder = ListBuilder::new(values_builder);
532
533        //  [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
534        builder.values().values().append_value(1);
535        builder.values().values().append_value(2);
536        builder.values().append(true);
537        builder.values().values().append_value(3);
538        builder.values().values().append_value(4);
539        builder.values().append(true);
540        builder.append(true);
541
542        builder.values().values().append_value(5);
543        builder.values().values().append_value(6);
544        builder.values().values().append_value(7);
545        builder.values().append(true);
546        builder.values().append(false);
547        builder.values().values().append_value(8);
548        builder.values().append(true);
549        builder.append(true);
550
551        builder.append(false);
552
553        builder.values().values().append_value(9);
554        builder.values().values().append_value(10);
555        builder.values().append(true);
556        builder.append(true);
557
558        let l1 = builder.finish();
559
560        assert_eq!(4, l1.len());
561        assert_eq!(1, l1.null_count());
562
563        assert_eq!(l1.value_offsets(), &[0, 2, 5, 5, 6]);
564        let l2 = l1.values().as_list::<i32>();
565
566        assert_eq!(6, l2.len());
567        assert_eq!(1, l2.null_count());
568        assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8, 10]);
569
570        let i1 = l2.values().as_primitive::<Int32Type>();
571        assert_eq!(10, i1.len());
572        assert_eq!(0, i1.null_count());
573        assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
574    }
575
576    #[test]
577    fn test_extend() {
578        let mut builder = ListBuilder::new(Int32Builder::new());
579        builder.extend([
580            Some(vec![Some(1), Some(2), Some(7), None]),
581            Some(vec![]),
582            Some(vec![Some(4), Some(5)]),
583            None,
584        ]);
585
586        let array = builder.finish();
587        assert_eq!(array.value_offsets(), [0, 4, 4, 6, 6]);
588        assert_eq!(array.null_count(), 1);
589        assert_eq!(array.logical_null_count(), 1);
590        assert!(array.is_null(3));
591        let elements = array.values().as_primitive::<Int32Type>();
592        assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]);
593        assert_eq!(elements.null_count(), 1);
594        assert_eq!(elements.logical_null_count(), 1);
595        assert!(elements.is_null(3));
596    }
597
598    #[test]
599    fn test_boxed_primitive_array_builder() {
600        let values_builder = make_builder(&DataType::Int32, 5);
601        let mut builder = ListBuilder::new(values_builder);
602
603        builder
604            .values()
605            .as_any_mut()
606            .downcast_mut::<Int32Builder>()
607            .expect("should be an Int32Builder")
608            .append_slice(&[1, 2, 3]);
609        builder.append(true);
610
611        builder
612            .values()
613            .as_any_mut()
614            .downcast_mut::<Int32Builder>()
615            .expect("should be an Int32Builder")
616            .append_slice(&[4, 5, 6]);
617        builder.append(true);
618
619        let arr = builder.finish();
620        assert_eq!(2, arr.len());
621
622        let elements = arr.values().as_primitive::<Int32Type>();
623        assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
624    }
625
626    #[test]
627    fn test_boxed_list_list_array_builder() {
628        // This test is same as `test_list_list_array_builder` but uses boxed builders.
629        let values_builder = make_builder(
630            &DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
631            10,
632        );
633        test_boxed_generic_list_generic_list_array_builder::<i32>(values_builder);
634    }
635
636    #[test]
637    fn test_boxed_large_list_large_list_array_builder() {
638        // This test is same as `test_list_list_array_builder` but uses boxed builders.
639        let values_builder = make_builder(
640            &DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int32, true))),
641            10,
642        );
643        test_boxed_generic_list_generic_list_array_builder::<i64>(values_builder);
644    }
645
646    fn test_boxed_generic_list_generic_list_array_builder<O: OffsetSizeTrait + PartialEq>(
647        values_builder: Box<dyn ArrayBuilder>,
648    ) {
649        let mut builder: GenericListBuilder<O, Box<dyn ArrayBuilder>> =
650            GenericListBuilder::<O, Box<dyn ArrayBuilder>>::new(values_builder);
651
652        //  [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
653        builder
654            .values()
655            .as_any_mut()
656            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
657            .expect("should be an (Large)ListBuilder")
658            .values()
659            .as_any_mut()
660            .downcast_mut::<Int32Builder>()
661            .expect("should be an Int32Builder")
662            .append_value(1);
663        builder
664            .values()
665            .as_any_mut()
666            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
667            .expect("should be an (Large)ListBuilder")
668            .values()
669            .as_any_mut()
670            .downcast_mut::<Int32Builder>()
671            .expect("should be an Int32Builder")
672            .append_value(2);
673        builder
674            .values()
675            .as_any_mut()
676            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
677            .expect("should be an (Large)ListBuilder")
678            .append(true);
679        builder
680            .values()
681            .as_any_mut()
682            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
683            .expect("should be an (Large)ListBuilder")
684            .values()
685            .as_any_mut()
686            .downcast_mut::<Int32Builder>()
687            .expect("should be an Int32Builder")
688            .append_value(3);
689        builder
690            .values()
691            .as_any_mut()
692            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
693            .expect("should be an (Large)ListBuilder")
694            .values()
695            .as_any_mut()
696            .downcast_mut::<Int32Builder>()
697            .expect("should be an Int32Builder")
698            .append_value(4);
699        builder
700            .values()
701            .as_any_mut()
702            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
703            .expect("should be an (Large)ListBuilder")
704            .append(true);
705        builder.append(true);
706
707        builder
708            .values()
709            .as_any_mut()
710            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
711            .expect("should be an (Large)ListBuilder")
712            .values()
713            .as_any_mut()
714            .downcast_mut::<Int32Builder>()
715            .expect("should be an Int32Builder")
716            .append_value(5);
717        builder
718            .values()
719            .as_any_mut()
720            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
721            .expect("should be an (Large)ListBuilder")
722            .values()
723            .as_any_mut()
724            .downcast_mut::<Int32Builder>()
725            .expect("should be an Int32Builder")
726            .append_value(6);
727        builder
728            .values()
729            .as_any_mut()
730            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
731            .expect("should be an (Large)ListBuilder")
732            .values()
733            .as_any_mut()
734            .downcast_mut::<Int32Builder>()
735            .expect("should be an (Large)ListBuilder")
736            .append_value(7);
737        builder
738            .values()
739            .as_any_mut()
740            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
741            .expect("should be an (Large)ListBuilder")
742            .append(true);
743        builder
744            .values()
745            .as_any_mut()
746            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
747            .expect("should be an (Large)ListBuilder")
748            .append(false);
749        builder
750            .values()
751            .as_any_mut()
752            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
753            .expect("should be an (Large)ListBuilder")
754            .values()
755            .as_any_mut()
756            .downcast_mut::<Int32Builder>()
757            .expect("should be an Int32Builder")
758            .append_value(8);
759        builder
760            .values()
761            .as_any_mut()
762            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
763            .expect("should be an (Large)ListBuilder")
764            .append(true);
765        builder.append(true);
766
767        builder.append(false);
768
769        builder
770            .values()
771            .as_any_mut()
772            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
773            .expect("should be an (Large)ListBuilder")
774            .values()
775            .as_any_mut()
776            .downcast_mut::<Int32Builder>()
777            .expect("should be an Int32Builder")
778            .append_value(9);
779        builder
780            .values()
781            .as_any_mut()
782            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
783            .expect("should be an (Large)ListBuilder")
784            .values()
785            .as_any_mut()
786            .downcast_mut::<Int32Builder>()
787            .expect("should be an Int32Builder")
788            .append_value(10);
789        builder
790            .values()
791            .as_any_mut()
792            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
793            .expect("should be an (Large)ListBuilder")
794            .append(true);
795        builder.append(true);
796
797        let l1 = builder.finish();
798
799        assert_eq!(4, l1.len());
800        assert_eq!(1, l1.null_count());
801
802        assert_eq!(l1.value_offsets(), &[0, 2, 5, 5, 6].map(O::usize_as));
803        let l2 = l1.values().as_list::<O>();
804
805        assert_eq!(6, l2.len());
806        assert_eq!(1, l2.null_count());
807        assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8, 10].map(O::usize_as));
808
809        let i1 = l2.values().as_primitive::<Int32Type>();
810        assert_eq!(10, i1.len());
811        assert_eq!(0, i1.null_count());
812        assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
813    }
814
815    #[test]
816    fn test_with_field() {
817        let field = Arc::new(Field::new("bar", DataType::Int32, false));
818        let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
819        builder.append_value([Some(1), Some(2), Some(3)]);
820        builder.append_null(); // This is fine as nullability refers to nullability of values
821        builder.append_value([Some(4)]);
822        let array = builder.finish();
823        assert_eq!(array.len(), 3);
824        assert_eq!(array.data_type(), &DataType::List(field.clone()));
825
826        builder.append_value([Some(4), Some(5)]);
827        let array = builder.finish();
828        assert_eq!(array.data_type(), &DataType::List(field));
829        assert_eq!(array.len(), 1);
830    }
831
832    #[test]
833    #[should_panic(expected = "Non-nullable field of ListArray \\\"item\\\" cannot contain nulls")]
834    fn test_checks_nullability() {
835        let field = Arc::new(Field::new_list_field(DataType::Int32, false));
836        let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
837        builder.append_value([Some(1), None]);
838        builder.finish();
839    }
840
841    #[test]
842    #[should_panic(expected = "ListArray expected data type Int64 got Int32")]
843    fn test_checks_data_type() {
844        let field = Arc::new(Field::new_list_field(DataType::Int64, false));
845        let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
846        builder.append_value([Some(1)]);
847        builder.finish();
848    }
849
850    #[test]
851    fn test_finish_preserve_values() {
852        let mut builder = ListBuilder::new(PreserveValuesMock::default());
853
854        builder.values().inner.append_value(1);
855        builder.append(true);
856
857        let arr = builder.finish_preserve_values();
858
859        assert_eq!(1, arr.len());
860        assert_eq!(1, builder.values().called);
861    }
862}
arrow_array/builder/generic_list_builder.rs

arrow_array/builder/
generic_list_builder.rs
