arrow_array/builder/
generic_list_builder.rs

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

arrow_array/builder/
generic_list_builder.rs
