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    pub fn finish(&mut self) -> GenericListArray<OffsetSize> {
303        let values = self.values_builder.finish();
304        let nulls = self.null_buffer_builder.finish();
305
306        let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
307        // Safety: Safe by construction
308        let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
309        self.offsets_builder.push(OffsetSize::zero());
310
311        let field = match &self.field {
312            Some(f) => f.clone(),
313            None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
314        };
315
316        GenericListArray::new(field, offsets, values, nulls)
317    }
318
319    /// Builds the [`GenericListArray`] without resetting the builder.
320    pub fn finish_cloned(&self) -> GenericListArray<OffsetSize> {
321        let values = self.values_builder.finish_cloned();
322        let nulls = self.null_buffer_builder.finish_cloned();
323
324        let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice());
325        // Safety: safe by construction
326        let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
327
328        let field = match &self.field {
329            Some(f) => f.clone(),
330            None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
331        };
332
333        GenericListArray::new(field, offsets, values, nulls)
334    }
335
336    fn finish_preserve_values(&mut self) -> GenericListArray<OffsetSize> {
337        let values = self.values_builder.finish_preserve_values();
338        let nulls = self.null_buffer_builder.finish();
339
340        let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
341        // Safety: Safe by construction
342        let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
343        self.offsets_builder.push(OffsetSize::zero());
344
345        let field = match &self.field {
346            Some(f) => f.clone(),
347            None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
348        };
349
350        GenericListArray::new(field, offsets, values, nulls)
351    }
352
353    /// Returns the current offsets buffer as a slice
354    pub fn offsets_slice(&self) -> &[OffsetSize] {
355        self.offsets_builder.as_slice()
356    }
357
358    /// Returns the current offsets buffer capacity, in offsets.
359    pub fn offsets_capacity(&self) -> usize {
360        self.offsets_builder.capacity()
361    }
362
363    /// Returns the current null buffer as a slice
364    pub fn validity_slice(&self) -> Option<&[u8]> {
365        self.null_buffer_builder.as_slice()
366    }
367
368    /// Returns the current null buffer allocated capacity, in bytes.
369    pub fn validity_capacity(&self) -> usize {
370        self.null_buffer_builder.allocated_size()
371    }
372}
373
374impl<O, B, V, E> Extend<Option<V>> for GenericListBuilder<O, B>
375where
376    O: OffsetSizeTrait,
377    B: ArrayBuilder + Extend<E>,
378    V: IntoIterator<Item = E>,
379{
380    #[inline]
381    fn extend<T: IntoIterator<Item = Option<V>>>(&mut self, iter: T) {
382        for v in iter {
383            match v {
384                Some(elements) => {
385                    self.values_builder.extend(elements);
386                    self.append(true);
387                }
388                None => self.append(false),
389            }
390        }
391    }
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397    use crate::Int32Array;
398    use crate::builder::{Int32Builder, ListBuilder, make_builder, tests::PreserveValuesMock};
399    use crate::cast::AsArray;
400    use crate::types::Int32Type;
401    use arrow_schema::DataType;
402
403    fn _test_generic_list_array_builder<O: OffsetSizeTrait>() {
404        let values_builder = Int32Builder::with_capacity(10);
405        let mut builder = GenericListBuilder::<O, _>::new(values_builder);
406
407        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
408        builder.values().append_value(0);
409        builder.values().append_value(1);
410        builder.values().append_value(2);
411        builder.append(true);
412        builder.values().append_value(3);
413        builder.values().append_value(4);
414        builder.values().append_value(5);
415        builder.append(true);
416        builder.values().append_value(6);
417        builder.values().append_value(7);
418        builder.append(true);
419        let list_array = builder.finish();
420
421        let list_values = list_array.values().as_primitive::<Int32Type>();
422        assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]);
423        assert_eq!(list_array.value_offsets(), [0, 3, 6, 8].map(O::usize_as));
424        assert_eq!(DataType::Int32, list_array.value_type());
425        assert_eq!(3, list_array.len());
426        assert_eq!(0, list_array.null_count());
427        assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]);
428        assert_eq!(O::from_usize(2).unwrap(), list_array.value_length(2));
429        for i in 0..3 {
430            assert!(list_array.is_valid(i));
431            assert!(!list_array.is_null(i));
432        }
433    }
434
435    #[test]
436    fn test_list_array_builder() {
437        _test_generic_list_array_builder::<i32>()
438    }
439
440    #[test]
441    fn test_large_list_array_builder() {
442        _test_generic_list_array_builder::<i64>()
443    }
444
445    fn _test_generic_list_array_builder_nulls<O: OffsetSizeTrait>() {
446        let values_builder = Int32Builder::with_capacity(10);
447        let mut builder = GenericListBuilder::<O, _>::new(values_builder);
448
449        //  [[0, 1, 2], null, [3, null, 5], [6, 7], null, null, [8]]
450        builder.values().append_value(0);
451        builder.values().append_value(1);
452        builder.values().append_value(2);
453        builder.append(true);
454        builder.append(false);
455        builder.values().append_value(3);
456        builder.values().append_null();
457        builder.values().append_value(5);
458        builder.append(true);
459        builder.values().append_value(6);
460        builder.values().append_value(7);
461        builder.append(true);
462        builder.append_nulls(2);
463        builder.values().append_value(8);
464        builder.append(true);
465
466        let list_array = builder.finish();
467
468        assert_eq!(DataType::Int32, list_array.value_type());
469        assert_eq!(7, list_array.len());
470        assert_eq!(3, list_array.null_count());
471        assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]);
472        assert_eq!(O::from_usize(9).unwrap(), list_array.value_offsets()[7]);
473        assert_eq!(O::from_usize(3).unwrap(), list_array.value_length(2));
474        assert!(list_array.is_null(4));
475        assert!(list_array.is_null(5));
476    }
477
478    #[test]
479    fn test_list_array_builder_nulls() {
480        _test_generic_list_array_builder_nulls::<i32>()
481    }
482
483    #[test]
484    fn test_large_list_array_builder_nulls() {
485        _test_generic_list_array_builder_nulls::<i64>()
486    }
487
488    #[test]
489    fn test_list_array_builder_finish() {
490        let values_builder = Int32Array::builder(5);
491        let mut builder = ListBuilder::new(values_builder);
492
493        builder.values().append_slice(&[1, 2, 3]);
494        builder.append(true);
495        builder.values().append_slice(&[4, 5, 6]);
496        builder.append(true);
497
498        let mut arr = builder.finish();
499        assert_eq!(2, arr.len());
500        assert!(builder.is_empty());
501
502        builder.values().append_slice(&[7, 8, 9]);
503        builder.append(true);
504        arr = builder.finish();
505        assert_eq!(1, arr.len());
506        assert!(builder.is_empty());
507    }
508
509    #[test]
510    fn test_list_array_builder_finish_cloned() {
511        let values_builder = Int32Array::builder(5);
512        let mut builder = ListBuilder::new(values_builder);
513
514        builder.values().append_slice(&[1, 2, 3]);
515        builder.append(true);
516        builder.values().append_slice(&[4, 5, 6]);
517        builder.append(true);
518
519        let mut arr = builder.finish_cloned();
520        assert_eq!(2, arr.len());
521        assert!(!builder.is_empty());
522
523        builder.values().append_slice(&[7, 8, 9]);
524        builder.append(true);
525        arr = builder.finish();
526        assert_eq!(3, arr.len());
527        assert!(builder.is_empty());
528    }
529
530    #[test]
531    fn test_list_list_array_builder() {
532        let primitive_builder = Int32Builder::with_capacity(10);
533        let values_builder = ListBuilder::new(primitive_builder);
534        let mut builder = ListBuilder::new(values_builder);
535
536        //  [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
537        builder.values().values().append_value(1);
538        builder.values().values().append_value(2);
539        builder.values().append(true);
540        builder.values().values().append_value(3);
541        builder.values().values().append_value(4);
542        builder.values().append(true);
543        builder.append(true);
544
545        builder.values().values().append_value(5);
546        builder.values().values().append_value(6);
547        builder.values().values().append_value(7);
548        builder.values().append(true);
549        builder.values().append(false);
550        builder.values().values().append_value(8);
551        builder.values().append(true);
552        builder.append(true);
553
554        builder.append(false);
555
556        builder.values().values().append_value(9);
557        builder.values().values().append_value(10);
558        builder.values().append(true);
559        builder.append(true);
560
561        let l1 = builder.finish();
562
563        assert_eq!(4, l1.len());
564        assert_eq!(1, l1.null_count());
565
566        assert_eq!(l1.value_offsets(), &[0, 2, 5, 5, 6]);
567        let l2 = l1.values().as_list::<i32>();
568
569        assert_eq!(6, l2.len());
570        assert_eq!(1, l2.null_count());
571        assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8, 10]);
572
573        let i1 = l2.values().as_primitive::<Int32Type>();
574        assert_eq!(10, i1.len());
575        assert_eq!(0, i1.null_count());
576        assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
577    }
578
579    #[test]
580    fn test_extend() {
581        let mut builder = ListBuilder::new(Int32Builder::new());
582        builder.extend([
583            Some(vec![Some(1), Some(2), Some(7), None]),
584            Some(vec![]),
585            Some(vec![Some(4), Some(5)]),
586            None,
587        ]);
588
589        let array = builder.finish();
590        assert_eq!(array.value_offsets(), [0, 4, 4, 6, 6]);
591        assert_eq!(array.null_count(), 1);
592        assert_eq!(array.logical_null_count(), 1);
593        assert!(array.is_null(3));
594        let elements = array.values().as_primitive::<Int32Type>();
595        assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]);
596        assert_eq!(elements.null_count(), 1);
597        assert_eq!(elements.logical_null_count(), 1);
598        assert!(elements.is_null(3));
599    }
600
601    #[test]
602    fn test_boxed_primitive_array_builder() {
603        let values_builder = make_builder(&DataType::Int32, 5);
604        let mut builder = ListBuilder::new(values_builder);
605
606        builder
607            .values()
608            .as_any_mut()
609            .downcast_mut::<Int32Builder>()
610            .expect("should be an Int32Builder")
611            .append_slice(&[1, 2, 3]);
612        builder.append(true);
613
614        builder
615            .values()
616            .as_any_mut()
617            .downcast_mut::<Int32Builder>()
618            .expect("should be an Int32Builder")
619            .append_slice(&[4, 5, 6]);
620        builder.append(true);
621
622        let arr = builder.finish();
623        assert_eq!(2, arr.len());
624
625        let elements = arr.values().as_primitive::<Int32Type>();
626        assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
627    }
628
629    #[test]
630    fn test_boxed_list_list_array_builder() {
631        // This test is same as `test_list_list_array_builder` but uses boxed builders.
632        let values_builder = make_builder(
633            &DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
634            10,
635        );
636        test_boxed_generic_list_generic_list_array_builder::<i32>(values_builder);
637    }
638
639    #[test]
640    fn test_boxed_large_list_large_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::LargeList(Arc::new(Field::new_list_field(DataType::Int32, true))),
644            10,
645        );
646        test_boxed_generic_list_generic_list_array_builder::<i64>(values_builder);
647    }
648
649    fn test_boxed_generic_list_generic_list_array_builder<O: OffsetSizeTrait + PartialEq>(
650        values_builder: Box<dyn ArrayBuilder>,
651    ) {
652        let mut builder: GenericListBuilder<O, Box<dyn ArrayBuilder>> =
653            GenericListBuilder::<O, Box<dyn ArrayBuilder>>::new(values_builder);
654
655        //  [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
656        builder
657            .values()
658            .as_any_mut()
659            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
660            .expect("should be an (Large)ListBuilder")
661            .values()
662            .as_any_mut()
663            .downcast_mut::<Int32Builder>()
664            .expect("should be an Int32Builder")
665            .append_value(1);
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(2);
676        builder
677            .values()
678            .as_any_mut()
679            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
680            .expect("should be an (Large)ListBuilder")
681            .append(true);
682        builder
683            .values()
684            .as_any_mut()
685            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
686            .expect("should be an (Large)ListBuilder")
687            .values()
688            .as_any_mut()
689            .downcast_mut::<Int32Builder>()
690            .expect("should be an Int32Builder")
691            .append_value(3);
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(4);
702        builder
703            .values()
704            .as_any_mut()
705            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
706            .expect("should be an (Large)ListBuilder")
707            .append(true);
708        builder.append(true);
709
710        builder
711            .values()
712            .as_any_mut()
713            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
714            .expect("should be an (Large)ListBuilder")
715            .values()
716            .as_any_mut()
717            .downcast_mut::<Int32Builder>()
718            .expect("should be an Int32Builder")
719            .append_value(5);
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(6);
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 (Large)ListBuilder")
739            .append_value(7);
740        builder
741            .values()
742            .as_any_mut()
743            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
744            .expect("should be an (Large)ListBuilder")
745            .append(true);
746        builder
747            .values()
748            .as_any_mut()
749            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
750            .expect("should be an (Large)ListBuilder")
751            .append(false);
752        builder
753            .values()
754            .as_any_mut()
755            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
756            .expect("should be an (Large)ListBuilder")
757            .values()
758            .as_any_mut()
759            .downcast_mut::<Int32Builder>()
760            .expect("should be an Int32Builder")
761            .append_value(8);
762        builder
763            .values()
764            .as_any_mut()
765            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
766            .expect("should be an (Large)ListBuilder")
767            .append(true);
768        builder.append(true);
769
770        builder.append(false);
771
772        builder
773            .values()
774            .as_any_mut()
775            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
776            .expect("should be an (Large)ListBuilder")
777            .values()
778            .as_any_mut()
779            .downcast_mut::<Int32Builder>()
780            .expect("should be an Int32Builder")
781            .append_value(9);
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(10);
792        builder
793            .values()
794            .as_any_mut()
795            .downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
796            .expect("should be an (Large)ListBuilder")
797            .append(true);
798        builder.append(true);
799
800        let l1 = builder.finish();
801
802        assert_eq!(4, l1.len());
803        assert_eq!(1, l1.null_count());
804
805        assert_eq!(l1.value_offsets(), &[0, 2, 5, 5, 6].map(O::usize_as));
806        let l2 = l1.values().as_list::<O>();
807
808        assert_eq!(6, l2.len());
809        assert_eq!(1, l2.null_count());
810        assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8, 10].map(O::usize_as));
811
812        let i1 = l2.values().as_primitive::<Int32Type>();
813        assert_eq!(10, i1.len());
814        assert_eq!(0, i1.null_count());
815        assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
816    }
817
818    #[test]
819    fn test_with_field() {
820        let field = Arc::new(Field::new("bar", DataType::Int32, false));
821        let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
822        builder.append_value([Some(1), Some(2), Some(3)]);
823        builder.append_null(); // This is fine as nullability refers to nullability of values
824        builder.append_value([Some(4)]);
825        let array = builder.finish();
826        assert_eq!(array.len(), 3);
827        assert_eq!(array.data_type(), &DataType::List(field.clone()));
828
829        builder.append_value([Some(4), Some(5)]);
830        let array = builder.finish();
831        assert_eq!(array.data_type(), &DataType::List(field));
832        assert_eq!(array.len(), 1);
833    }
834
835    #[test]
836    #[should_panic(expected = "Non-nullable field of ListArray \\\"item\\\" cannot contain nulls")]
837    fn test_checks_nullability() {
838        let field = Arc::new(Field::new_list_field(DataType::Int32, false));
839        let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
840        builder.append_value([Some(1), None]);
841        builder.finish();
842    }
843
844    #[test]
845    #[should_panic(expected = "ListArray expected data type Int64 got Int32")]
846    fn test_checks_data_type() {
847        let field = Arc::new(Field::new_list_field(DataType::Int64, false));
848        let mut builder = ListBuilder::new(Int32Builder::new()).with_field(field.clone());
849        builder.append_value([Some(1)]);
850        builder.finish();
851    }
852
853    #[test]
854    fn test_finish_preserve_values() {
855        let mut builder = ListBuilder::new(PreserveValuesMock::default());
856
857        builder.values().inner.append_value(1);
858        builder.append(true);
859
860        let arr = builder.finish_preserve_values();
861
862        assert_eq!(1, arr.len());
863        assert_eq!(1, builder.values().called);
864    }
865}
arrow_array/builder/generic_list_builder.rs

arrow_array/builder/
generic_list_builder.rs
