arrow_array/array/
list_array.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::array::{get_offsets_from_buffer, make_array, print_long_array};
6use crate::builder::{ArrayBuilder, GenericListBuilder, PrimitiveBuilder};
7use crate::{
8    Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, FixedSizeListArray,
9    iterator::GenericListArrayIter, new_empty_array,
10};
11use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer};
12use arrow_data::{ArrayData, ArrayDataBuilder};
13use arrow_schema::{ArrowError, DataType, FieldRef};
14use num_integer::Integer;
15use std::any::Any;
16use std::sync::Arc;
17
18/// A type that can be used within a variable-size array to encode offset information
19///
20/// See [`ListArray`], [`LargeListArray`], [`BinaryArray`], [`LargeBinaryArray`],
21/// [`StringArray`] and [`LargeStringArray`]
22///
23/// [`BinaryArray`]: crate::array::BinaryArray
24/// [`LargeBinaryArray`]: crate::array::LargeBinaryArray
25/// [`StringArray`]: crate::array::StringArray
26/// [`LargeStringArray`]: crate::array::LargeStringArray
27pub trait OffsetSizeTrait:
28    ArrowNativeType + std::ops::AddAssign + Integer + num_traits::CheckedAdd + num_traits::CheckedSub
29{
30    /// True for 64 bit offset size and false for 32 bit offset size
31    const IS_LARGE: bool;
32    /// Prefix for the offset size
33    const PREFIX: &'static str;
34    /// The max `usize` offset
35    const MAX_OFFSET: usize;
36}
37
38impl OffsetSizeTrait for i32 {
39    const IS_LARGE: bool = false;
40    const PREFIX: &'static str = "";
41    const MAX_OFFSET: usize = i32::MAX as usize;
42}
43
44impl OffsetSizeTrait for i64 {
45    const IS_LARGE: bool = true;
46    const PREFIX: &'static str = "Large";
47    const MAX_OFFSET: usize = i64::MAX as usize;
48}
49
50/// An array of [variable length lists], similar to JSON arrays
51/// (e.g. `["A", "B", "C"]`). This struct specifically represents
52/// the [list layout]. Refer to [`GenericListViewArray`] for the
53/// [list-view layout].
54///
55/// Lists are represented using `offsets` into a `values` child
56/// array. Offsets are stored in two adjacent entries of an
57/// [`OffsetBuffer`].
58///
59/// Arrow defines [`ListArray`] with `i32` offsets and
60/// [`LargeListArray`] with `i64` offsets.
61///
62/// Use [`GenericListBuilder`] to construct a [`GenericListArray`].
63///
64/// # Representation
65///
66/// A [`ListArray`] can represent a list of values of any other
67/// supported Arrow type. Each element of the `ListArray` itself is
68/// a list which may be empty, may contain NULL and non-null values,
69/// or may itself be NULL.
70///
71/// For example, the `ListArray` shown in the following diagram stores
72/// lists of strings. Note that `[]` represents an empty (length
73/// 0), but non NULL list.
74///
75/// ```text
76/// ┌─────────────┐
77/// │   [A,B,C]   │
78/// ├─────────────┤
79/// │     []      │
80/// ├─────────────┤
81/// │    NULL     │
82/// ├─────────────┤
83/// │     [D]     │
84/// ├─────────────┤
85/// │  [NULL, F]  │
86/// └─────────────┘
87/// ```
88///
89/// The `values` are stored in a child [`StringArray`] and the offsets
90/// are stored in an [`OffsetBuffer`] as shown in the following
91/// diagram. The logical values and offsets are shown on the left, and
92/// the actual `ListArray` encoding on the right.
93///
94/// ```text
95///                                         ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
96///                                                                 ┌ ─ ─ ─ ─ ─ ─ ┐    │
97///  ┌─────────────┐  ┌───────┐             │     ┌───┐   ┌───┐       ┌───┐ ┌───┐
98///  │   [A,B,C]   │  │ (0,3) │                   │ 1 │   │ 0 │     │ │ 1 │ │ A │ │ 0  │
99///  ├─────────────┤  ├───────┤             │     ├───┤   ├───┤       ├───┤ ├───┤
100///  │ [] (empty)  │  │ (3,3) │                   │ 1 │   │ 3 │     │ │ 1 │ │ B │ │ 1  │
101///  ├─────────────┤  ├───────┤             │     ├───┤   ├───┤       ├───┤ ├───┤
102///  │    NULL     │  │ (3,3) │                   │ 0 │   │ 3 │     │ │ 1 │ │ C │ │ 2  │
103///  ├─────────────┤  ├───────┤             │     ├───┤   ├───┤       ├───┤ ├───┤
104///  │     [D]     │  │ (3,4) │                   │ 1 │   │ 3 │     │ │ 1 │ │ D │ │ 3  │
105///  ├─────────────┤  ├───────┤             │     ├───┤   ├───┤       ├───┤ ├───┤
106///  │  [NULL, F]  │  │ (4,6) │                   │ 1 │   │ 4 │     │ │ 0 │ │ ? │ │ 4  │
107///  └─────────────┘  └───────┘             │     └───┘   ├───┤       ├───┤ ├───┤
108///                                                       │ 6 │     │ │ 1 │ │ F │ │ 5  │
109///                                         │  Validity   └───┘       └───┘ └───┘
110///     Logical       Logical                  (nulls)   Offsets    │    Values   │    │
111///      Values       Offsets               │                           (Array)
112///                                                                 └ ─ ─ ─ ─ ─ ─ ┘    │
113///                 (offsets[i],            │   ListArray
114///                offsets[i+1])                                                       │
115///                                         └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
116/// ```
117///
118/// # Slicing
119///
120/// Slicing a `ListArray` creates a new `ListArray` without copying any data,
121/// but this means the [`Self::values`] and [`Self::offsets`] may have "unused" data
122///
123/// For example, calling `slice(1, 3)` on the `ListArray` in the above example
124/// would result in the following. Note
125///
126/// 1. `Values` array is unchanged
127/// 2. `Offsets` do not start at `0`, nor cover all values in the Values array.
128///
129/// ```text
130///                                 ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
131///                                                         ┌ ─ ─ ─ ─ ─ ─ ┐    │  ╔═══╗
132///                                 │                         ╔═══╗ ╔═══╗         ║   ║  Not used
133///                                                         │ ║ 1 ║ ║ A ║ │ 0  │  ╚═══╝
134///  ┌─────────────┐  ┌───────┐     │     ┌───┐   ┌───┐       ╠═══╣ ╠═══╣
135///  │ [] (empty)  │  │ (3,3) │           │ 1 │   │ 3 │     │ ║ 1 ║ ║ B ║ │ 1  │
136///  ├─────────────┤  ├───────┤     │     ├───┤   ├───┤       ╠═══╣ ╠═══╣
137///  │    NULL     │  │ (3,3) │           │ 0 │   │ 3 │     │ ║ 1 ║ ║ C ║ │ 2  │
138///  ├─────────────┤  ├───────┤     │     ├───┤   ├───┤       ╚═══╝ ╚═══╝
139///  │     [D]     │  │ (3,4) │           │ 1 │   │ 3 │     │ │ 1 │ │ D │ │ 3  │
140///  └─────────────┘  └───────┘     │     └───┘   ├───┤       ╔═══╗ ╔═══╗
141///                                               │ 4 │     │ ║ 0 ║ ║ ? ║ │ 4  │
142///                                 │             └───┘       ╠═══╣ ╠═══╣
143///                                                         │ ║ 1 ║ ║ F ║ │ 5  │
144///                                 │  Validity               ╚═══╝ ╚═══╝
145///     Logical       Logical          (nulls)   Offsets    │    Values   │    │
146///      Values       Offsets       │                           (Array)
147///                                                         └ ─ ─ ─ ─ ─ ─ ┘    │
148///                 (offsets[i],    │   ListArray
149///                offsets[i+1])                                               │
150///                                 └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
151/// ```
152///
153/// [`StringArray`]: crate::array::StringArray
154/// [`GenericListViewArray`]: crate::array::GenericListViewArray
155/// [variable length lists]: https://arrow.apache.org/docs/format/Columnar.html#variable-size-list-layout
156/// [list layout]: https://arrow.apache.org/docs/format/Columnar.html#list-layout
157/// [list-view layout]: https://arrow.apache.org/docs/format/Columnar.html#listview-layout
158pub struct GenericListArray<OffsetSize: OffsetSizeTrait> {
159    data_type: DataType,
160    nulls: Option<NullBuffer>,
161    values: ArrayRef,
162    value_offsets: OffsetBuffer<OffsetSize>,
163}
164
165impl<OffsetSize: OffsetSizeTrait> Clone for GenericListArray<OffsetSize> {
166    fn clone(&self) -> Self {
167        Self {
168            data_type: self.data_type.clone(),
169            nulls: self.nulls.clone(),
170            values: self.values.clone(),
171            value_offsets: self.value_offsets.clone(),
172        }
173    }
174}
175
176impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
177    /// The data type constructor of list array.
178    /// The input is the schema of the child array and
179    /// the output is the [`DataType`], List or LargeList.
180    pub const DATA_TYPE_CONSTRUCTOR: fn(FieldRef) -> DataType = if OffsetSize::IS_LARGE {
181        DataType::LargeList
182    } else {
183        DataType::List
184    };
185
186    /// Create a new [`GenericListArray`] from the provided parts
187    ///
188    /// # Errors
189    ///
190    /// Errors if
191    ///
192    /// * `offsets.len() - 1 != nulls.len()`
193    /// * `offsets.last() > values.len()`
194    /// * `!field.is_nullable() && values.logical_null_count() != 0`
195    /// * `field.data_type() != values.data_type()`
196    pub fn try_new(
197        field: FieldRef,
198        offsets: OffsetBuffer<OffsetSize>,
199        values: ArrayRef,
200        nulls: Option<NullBuffer>,
201    ) -> Result<Self, ArrowError> {
202        let len = offsets.len() - 1; // Offsets guaranteed to not be empty
203        let end_offset = offsets.last().as_usize();
204        // don't need to check other values of `offsets` because they are checked
205        // during construction of `OffsetBuffer`
206        if end_offset > values.len() {
207            return Err(ArrowError::InvalidArgumentError(format!(
208                "Max offset of {end_offset} exceeds length of values {}",
209                values.len()
210            )));
211        }
212
213        if let Some(n) = nulls.as_ref()
214            && n.len() != len
215        {
216            return Err(ArrowError::InvalidArgumentError(format!(
217                "Incorrect length of null buffer for {}ListArray, expected {len} got {}",
218                OffsetSize::PREFIX,
219                n.len(),
220            )));
221        }
222        if !field.is_nullable() && values.logical_null_count() != 0 {
223            return Err(ArrowError::InvalidArgumentError(format!(
224                "Non-nullable field of {}ListArray {:?} cannot contain nulls",
225                OffsetSize::PREFIX,
226                field.name()
227            )));
228        }
229
230        if field.data_type() != values.data_type() {
231            return Err(ArrowError::InvalidArgumentError(format!(
232                "{}ListArray expected data type {} got {} for {:?}",
233                OffsetSize::PREFIX,
234                field.data_type(),
235                values.data_type(),
236                field.name()
237            )));
238        }
239
240        Ok(Self {
241            data_type: Self::DATA_TYPE_CONSTRUCTOR(field),
242            nulls,
243            values,
244            value_offsets: offsets,
245        })
246    }
247
248    /// Create a new [`GenericListArray`] from the provided parts
249    ///
250    /// # Panics
251    ///
252    /// Panics if [`Self::try_new`] returns an error
253    pub fn new(
254        field: FieldRef,
255        offsets: OffsetBuffer<OffsetSize>,
256        values: ArrayRef,
257        nulls: Option<NullBuffer>,
258    ) -> Self {
259        Self::try_new(field, offsets, values, nulls).unwrap()
260    }
261
262    /// Create a new [`GenericListArray`] from the provided parts without validation.
263    ///
264    /// # Safety
265    /// - `offsets.len() - 1 == nulls.len()` if `nulls` is `Some`
266    /// - `offsets.last() <= values.len()`
267    /// - `field.data_type() == values.data_type()`
268    pub unsafe fn new_unchecked(
269        field: FieldRef,
270        offsets: OffsetBuffer<OffsetSize>,
271        values: ArrayRef,
272        nulls: Option<NullBuffer>,
273    ) -> Self {
274        if cfg!(feature = "force_validate") {
275            return Self::new(field, offsets, values, nulls);
276        }
277        Self {
278            data_type: Self::DATA_TYPE_CONSTRUCTOR(field),
279            nulls,
280            values,
281            value_offsets: offsets,
282        }
283    }
284
285    /// Create a new [`GenericListArray`] of length `len` where all values are null
286    pub fn new_null(field: FieldRef, len: usize) -> Self {
287        let values = new_empty_array(field.data_type());
288        Self {
289            data_type: Self::DATA_TYPE_CONSTRUCTOR(field),
290            nulls: Some(NullBuffer::new_null(len)),
291            value_offsets: OffsetBuffer::new_zeroed(len),
292            values,
293        }
294    }
295
296    /// Deconstruct this array into its constituent parts
297    pub fn into_parts(
298        self,
299    ) -> (
300        FieldRef,
301        OffsetBuffer<OffsetSize>,
302        ArrayRef,
303        Option<NullBuffer>,
304    ) {
305        let (DataType::List(f) | DataType::LargeList(f)) = self.data_type else {
306            unreachable!()
307        };
308        (f, self.value_offsets, self.values, self.nulls)
309    }
310
311    /// The field that describes the values of this list.
312    pub fn value_field(&self) -> &FieldRef {
313        match &self.data_type {
314            DataType::List(f) | DataType::LargeList(f) => f,
315            _ => unreachable!(),
316        }
317    }
318
319    /// Returns a reference to the offsets of this list
320    ///
321    /// Unlike [`Self::value_offsets`] this returns the [`OffsetBuffer`]
322    /// allowing for zero-copy cloning.
323    ///
324    /// Notes: The `offsets` may not start at 0 and may not cover all values in
325    /// [`Self::values`]. This can happen when the list array was sliced via
326    /// [`Self::slice`]. See documentation for [`Self`] for more details.
327    #[inline]
328    pub fn offsets(&self) -> &OffsetBuffer<OffsetSize> {
329        &self.value_offsets
330    }
331
332    /// Returns a reference to the values of this list
333    ///
334    /// Note: The list array may not refer to all values in the `values` array.
335    /// For example if the list array was sliced via [`Self::slice`] values will
336    /// still contain values both before and after the slice. See documentation
337    /// for [`Self`] for more details.
338    #[inline]
339    pub fn values(&self) -> &ArrayRef {
340        &self.values
341    }
342
343    /// Returns a clone of the value type of this list.
344    pub fn value_type(&self) -> DataType {
345        self.values.data_type().clone()
346    }
347
348    /// Returns ith value of this list array.
349    ///
350    /// Note: This method does not check for nulls and the value is arbitrary
351    /// if [`is_null`](Self::is_null) returns true for the index.
352    ///
353    /// # Safety
354    /// Caller must ensure that the index is within the array bounds
355    pub unsafe fn value_unchecked(&self, i: usize) -> ArrayRef {
356        let end = unsafe { self.value_offsets().get_unchecked(i + 1).as_usize() };
357        let start = unsafe { self.value_offsets().get_unchecked(i).as_usize() };
358        self.values.slice(start, end - start)
359    }
360
361    /// Returns ith value of this list array.
362    ///
363    /// Note: This method does not check for nulls and the value is arbitrary
364    /// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
365    ///
366    /// # Panics
367    /// Panics if index `i` is out of bounds
368    pub fn value(&self, i: usize) -> ArrayRef {
369        let end = self.value_offsets()[i + 1].as_usize();
370        let start = self.value_offsets()[i].as_usize();
371        self.values.slice(start, end - start)
372    }
373
374    /// Returns the offset values in the offsets buffer.
375    ///
376    /// See [`Self::offsets`] for more details.
377    #[inline]
378    pub fn value_offsets(&self) -> &[OffsetSize] {
379        &self.value_offsets
380    }
381
382    /// Returns the length for value at index `i`.
383    ///
384    /// # Panics
385    /// Panics if `i >= self.len()`
386    #[inline]
387    pub fn value_length(&self, i: usize) -> OffsetSize {
388        let offsets = self.value_offsets();
389        offsets[i + 1] - offsets[i]
390    }
391
392    /// constructs a new iterator
393    pub fn iter<'a>(&'a self) -> GenericListArrayIter<'a, OffsetSize> {
394        GenericListArrayIter::<'a, OffsetSize>::new(self)
395    }
396
397    #[inline]
398    fn get_type(data_type: &DataType) -> Option<&DataType> {
399        match (OffsetSize::IS_LARGE, data_type) {
400            (true, DataType::LargeList(child)) | (false, DataType::List(child)) => {
401                Some(child.data_type())
402            }
403            _ => None,
404        }
405    }
406
407    /// Returns a zero-copy slice of this array with the indicated offset and length.
408    ///
409    /// Notes: this method does *NOT* slice the underlying values array or modify
410    /// the values in the offsets buffer. See [`Self::values`] and
411    /// [`Self::offsets`] for more information.
412    ///
413    /// # Panics
414    /// Panics if `offset + length > self.len()`
415    pub fn slice(&self, offset: usize, length: usize) -> Self {
416        Self {
417            data_type: self.data_type.clone(),
418            nulls: self.nulls.as_ref().map(|n| n.slice(offset, length)),
419            values: self.values.clone(),
420            value_offsets: self.value_offsets.slice(offset, length),
421        }
422    }
423
424    /// Creates a [`GenericListArray`] from an iterator of primitive values
425    /// # Example
426    /// ```
427    /// # use arrow_array::ListArray;
428    /// # use arrow_array::types::Int32Type;
429    ///
430    /// let data = vec![
431    ///    Some(vec![Some(0), Some(1), Some(2)]),
432    ///    None,
433    ///    Some(vec![Some(3), None, Some(5)]),
434    ///    Some(vec![Some(6), Some(7)]),
435    /// ];
436    /// let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
437    /// println!("{:?}", list_array);
438    /// ```
439    pub fn from_iter_primitive<T, P, I>(iter: I) -> Self
440    where
441        T: ArrowPrimitiveType,
442        P: IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
443        I: IntoIterator<Item = Option<P>>,
444    {
445        Self::from_nested_iter::<PrimitiveBuilder<T>, T::Native, P, I>(iter)
446    }
447
448    /// Creates a [`GenericListArray`] from a nested iterator of values.
449    /// This method works for any values type that has a corresponding builder that implements the
450    /// `Extend` trait. That includes all numeric types, booleans, binary and string types and also
451    /// dictionary encoded binary and strings.
452    ///
453    /// # Example
454    /// ```
455    /// # use arrow_array::ListArray;
456    /// # use arrow_array::types::Int32Type;
457    /// # use arrow_array::builder::StringDictionaryBuilder;
458    /// let data = vec![
459    ///    Some(vec![Some("foo"), Some("bar"), Some("baz")]),
460    ///    None,
461    ///    Some(vec![Some("bar"), None, Some("foo")]),
462    ///    Some(vec![]),
463    /// ];
464    /// let list_array = ListArray::from_nested_iter::<StringDictionaryBuilder<Int32Type>, _, _, _>(data);
465    /// println!("{:?}", list_array);
466    /// ```
467    pub fn from_nested_iter<B, T, P, I>(iter: I) -> Self
468    where
469        B: ArrayBuilder + Default + Extend<Option<T>>,
470        P: IntoIterator<Item = Option<T>>,
471        I: IntoIterator<Item = Option<P>>,
472    {
473        let iter = iter.into_iter();
474        let size_hint = iter.size_hint().0;
475        let mut builder = GenericListBuilder::with_capacity(B::default(), size_hint);
476
477        for i in iter {
478            match i {
479                Some(p) => {
480                    builder.values().extend(p);
481                    builder.append(true);
482                }
483                None => builder.append(false),
484            }
485        }
486        builder.finish()
487    }
488}
489
490impl<OffsetSize: OffsetSizeTrait> From<ArrayData> for GenericListArray<OffsetSize> {
491    fn from(data: ArrayData) -> Self {
492        Self::try_new_from_array_data(data)
493            .expect("Expected infallible creation of GenericListArray from ArrayDataRef failed")
494    }
495}
496
497impl<OffsetSize: OffsetSizeTrait> From<GenericListArray<OffsetSize>> for ArrayData {
498    fn from(array: GenericListArray<OffsetSize>) -> Self {
499        let len = array.len();
500        let builder = ArrayDataBuilder::new(array.data_type)
501            .len(len)
502            .nulls(array.nulls)
503            .buffers(vec![array.value_offsets.into_inner().into_inner()])
504            .child_data(vec![array.values.to_data()]);
505
506        unsafe { builder.build_unchecked() }
507    }
508}
509
510impl<OffsetSize: OffsetSizeTrait> From<FixedSizeListArray> for GenericListArray<OffsetSize> {
511    fn from(value: FixedSizeListArray) -> Self {
512        let (field, size) = match value.data_type() {
513            DataType::FixedSizeList(f, size) => (f, *size as usize),
514            _ => unreachable!(),
515        };
516
517        let offsets = OffsetBuffer::from_repeated_length(size, value.len());
518
519        Self {
520            data_type: Self::DATA_TYPE_CONSTRUCTOR(field.clone()),
521            nulls: value.nulls().cloned(),
522            values: value.values().clone(),
523            value_offsets: offsets,
524        }
525    }
526}
527
528impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
529    fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
530        let (data_type, len, nulls, offset, mut buffers, mut child_data) = data.into_parts();
531
532        if buffers.len() != 1 {
533            return Err(ArrowError::InvalidArgumentError(format!(
534                "ListArray data should contain a single buffer only (value offsets), had {}",
535                buffers.len()
536            )));
537        }
538        let buffer = buffers.pop().expect("checked above");
539
540        if child_data.len() != 1 {
541            return Err(ArrowError::InvalidArgumentError(format!(
542                "ListArray should contain a single child array (values array), had {}",
543                child_data.len()
544            )));
545        }
546
547        let values = child_data.pop().expect("checked above");
548
549        if let Some(child_data_type) = Self::get_type(&data_type) {
550            if values.data_type() != child_data_type {
551                return Err(ArrowError::InvalidArgumentError(format!(
552                    "[Large]ListArray's child datatype {:?} does not \
553                             correspond to the List's datatype {:?}",
554                    values.data_type(),
555                    child_data_type
556                )));
557            }
558        } else {
559            return Err(ArrowError::InvalidArgumentError(format!(
560                "[Large]ListArray's datatype must be [Large]ListArray(). It is {data_type:?}",
561            )));
562        }
563
564        let values = make_array(values);
565        // SAFETY:
566        // ArrayData is valid, and verified type above
567        let value_offsets = unsafe { get_offsets_from_buffer(buffer, offset, len) };
568
569        Ok(Self {
570            data_type,
571            nulls,
572            values,
573            value_offsets,
574        })
575    }
576}
577
578/// SAFETY: Correctly implements the contract of Arrow Arrays
579unsafe impl<OffsetSize: OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
580    fn as_any(&self) -> &dyn Any {
581        self
582    }
583
584    fn to_data(&self) -> ArrayData {
585        self.clone().into()
586    }
587
588    fn into_data(self) -> ArrayData {
589        self.into()
590    }
591
592    fn data_type(&self) -> &DataType {
593        &self.data_type
594    }
595
596    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
597        Arc::new(self.slice(offset, length))
598    }
599
600    fn len(&self) -> usize {
601        self.value_offsets.len() - 1
602    }
603
604    fn is_empty(&self) -> bool {
605        self.value_offsets.len() <= 1
606    }
607
608    fn shrink_to_fit(&mut self) {
609        if let Some(nulls) = &mut self.nulls {
610            nulls.shrink_to_fit();
611        }
612        self.values.shrink_to_fit();
613        self.value_offsets.shrink_to_fit();
614    }
615
616    fn offset(&self) -> usize {
617        0
618    }
619
620    fn nulls(&self) -> Option<&NullBuffer> {
621        self.nulls.as_ref()
622    }
623
624    fn logical_null_count(&self) -> usize {
625        // More efficient that the default implementation
626        self.null_count()
627    }
628
629    fn get_buffer_memory_size(&self) -> usize {
630        let mut size = self.values.get_buffer_memory_size();
631        size += self.value_offsets.inner().inner().capacity();
632        if let Some(n) = self.nulls.as_ref() {
633            size += n.buffer().capacity();
634        }
635        size
636    }
637
638    fn get_array_memory_size(&self) -> usize {
639        let mut size = std::mem::size_of::<Self>() + self.values.get_array_memory_size();
640        size += self.value_offsets.inner().inner().capacity();
641        if let Some(n) = self.nulls.as_ref() {
642            size += n.buffer().capacity();
643        }
644        size
645    }
646
647    #[cfg(feature = "pool")]
648    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
649        self.value_offsets.claim(pool);
650        self.values.claim(pool);
651        if let Some(nulls) = &self.nulls {
652            nulls.claim(pool);
653        }
654    }
655}
656
657impl<OffsetSize: OffsetSizeTrait> super::ListLikeArray for GenericListArray<OffsetSize> {
658    fn values(&self) -> &ArrayRef {
659        self.values()
660    }
661
662    fn element_range(&self, index: usize) -> std::ops::Range<usize> {
663        let offsets = self.offsets();
664        let start = offsets[index].as_usize();
665        let end = offsets[index + 1].as_usize();
666        start..end
667    }
668}
669
670impl<'a, OffsetSize: OffsetSizeTrait> IntoIterator for &'a GenericListArray<OffsetSize> {
671    type Item = Option<ArrayRef>;
672    type IntoIter = GenericListArrayIter<'a, OffsetSize>;
673
674    fn into_iter(self) -> Self::IntoIter {
675        GenericListArrayIter::<'a, OffsetSize>::new(self)
676    }
677}
678
679impl<OffsetSize: OffsetSizeTrait> ArrayAccessor for &GenericListArray<OffsetSize> {
680    type Item = ArrayRef;
681
682    fn value(&self, index: usize) -> Self::Item {
683        GenericListArray::value(self, index)
684    }
685
686    unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
687        GenericListArray::value(self, index)
688    }
689}
690
691impl<OffsetSize: OffsetSizeTrait> std::fmt::Debug for GenericListArray<OffsetSize> {
692    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
693        let prefix = OffsetSize::PREFIX;
694
695        write!(f, "{prefix}ListArray\n[\n")?;
696        print_long_array(self, f, &mut |index, f| {
697            std::fmt::Debug::fmt(&self.value(index), f)
698        })?;
699        write!(f, "]")
700    }
701}
702
703/// A [`GenericListArray`] of variable size lists, storing offsets as `i32`.
704///
705/// See [`ListBuilder`](crate::builder::ListBuilder) for how to construct a [`ListArray`]
706pub type ListArray = GenericListArray<i32>;
707
708/// A [`GenericListArray`] of variable size lists, storing offsets as `i64`.
709///
710/// See [`LargeListBuilder`](crate::builder::LargeListBuilder) for how to construct a [`LargeListArray`]
711pub type LargeListArray = GenericListArray<i64>;
712
713#[cfg(test)]
714mod tests {
715    use super::*;
716    use crate::builder::{
717        BooleanBuilder, FixedSizeListBuilder, Int32Builder, ListBuilder, StringBuilder,
718        StringDictionaryBuilder, UnionBuilder,
719    };
720    use crate::cast::AsArray;
721    use crate::types::{Int8Type, Int32Type};
722    use crate::{
723        BooleanArray, Int8Array, Int8DictionaryArray, Int32Array, Int64Array, StringArray,
724    };
725    use arrow_buffer::{Buffer, ScalarBuffer, bit_util};
726    use arrow_schema::Field;
727
728    fn create_from_buffers() -> ListArray {
729        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
730        let values = Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7]);
731        let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 6, 8]));
732        let field = Arc::new(Field::new_list_field(DataType::Int32, true));
733        ListArray::new(field, offsets, Arc::new(values), None)
734    }
735
736    #[test]
737    fn test_from_iter_primitive() {
738        let data = vec![
739            Some(vec![Some(0), Some(1), Some(2)]),
740            Some(vec![Some(3), Some(4), Some(5)]),
741            Some(vec![Some(6), Some(7)]),
742        ];
743        let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
744
745        let another = create_from_buffers();
746        assert_eq!(list_array, another)
747    }
748
749    #[test]
750    fn test_empty_list_array() {
751        // Construct an empty value array
752        let value_data = ArrayData::builder(DataType::Int32)
753            .len(0)
754            .add_buffer(Buffer::from([]))
755            .build()
756            .unwrap();
757
758        // Construct an empty offset buffer
759        let value_offsets = Buffer::from([]);
760
761        // Construct a list array from the above two
762        let list_data_type =
763            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
764        let list_data = ArrayData::builder(list_data_type)
765            .len(0)
766            .add_buffer(value_offsets)
767            .add_child_data(value_data)
768            .build()
769            .unwrap();
770
771        let list_array = ListArray::from(list_data);
772        assert_eq!(list_array.len(), 0)
773    }
774
775    #[test]
776    fn test_list_array() {
777        // Construct a value array
778        let value_data = ArrayData::builder(DataType::Int32)
779            .len(8)
780            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
781            .build()
782            .unwrap();
783
784        // Construct a buffer for value offsets, for the nested array:
785        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
786        let value_offsets = Buffer::from_slice_ref([0, 3, 6, 8]);
787
788        // Construct a list array from the above two
789        let list_data_type =
790            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
791        let list_data = ArrayData::builder(list_data_type.clone())
792            .len(3)
793            .add_buffer(value_offsets.clone())
794            .add_child_data(value_data.clone())
795            .build()
796            .unwrap();
797        let list_array = ListArray::from(list_data);
798
799        let values = list_array.values();
800        assert_eq!(value_data, values.to_data());
801        assert_eq!(DataType::Int32, list_array.value_type());
802        assert_eq!(3, list_array.len());
803        assert_eq!(0, list_array.null_count());
804        assert_eq!(6, list_array.value_offsets()[2]);
805        assert_eq!(2, list_array.value_length(2));
806        assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0));
807        assert_eq!(
808            0,
809            unsafe { list_array.value_unchecked(0) }
810                .as_primitive::<Int32Type>()
811                .value(0)
812        );
813        for i in 0..3 {
814            assert!(list_array.is_valid(i));
815            assert!(!list_array.is_null(i));
816        }
817
818        // Now test with a non-zero offset (skip first element)
819        //  [[3, 4, 5], [6, 7]]
820        let list_data = ArrayData::builder(list_data_type)
821            .len(2)
822            .offset(1)
823            .add_buffer(value_offsets)
824            .add_child_data(value_data.clone())
825            .build()
826            .unwrap();
827        let list_array = ListArray::from(list_data);
828
829        let values = list_array.values();
830        assert_eq!(value_data, values.to_data());
831        assert_eq!(DataType::Int32, list_array.value_type());
832        assert_eq!(2, list_array.len());
833        assert_eq!(0, list_array.null_count());
834        assert_eq!(6, list_array.value_offsets()[1]);
835        assert_eq!(2, list_array.value_length(1));
836        assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0));
837        assert_eq!(
838            3,
839            unsafe { list_array.value_unchecked(0) }
840                .as_primitive::<Int32Type>()
841                .value(0)
842        );
843    }
844
845    #[test]
846    fn test_large_list_array() {
847        // Construct a value array
848        let value_data = ArrayData::builder(DataType::Int32)
849            .len(8)
850            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
851            .build()
852            .unwrap();
853
854        // Construct a buffer for value offsets, for the nested array:
855        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
856        let value_offsets = Buffer::from_slice_ref([0i64, 3, 6, 8]);
857
858        // Construct a list array from the above two
859        let list_data_type = DataType::new_large_list(DataType::Int32, false);
860        let list_data = ArrayData::builder(list_data_type.clone())
861            .len(3)
862            .add_buffer(value_offsets.clone())
863            .add_child_data(value_data.clone())
864            .build()
865            .unwrap();
866        let list_array = LargeListArray::from(list_data);
867
868        let values = list_array.values();
869        assert_eq!(value_data, values.to_data());
870        assert_eq!(DataType::Int32, list_array.value_type());
871        assert_eq!(3, list_array.len());
872        assert_eq!(0, list_array.null_count());
873        assert_eq!(6, list_array.value_offsets()[2]);
874        assert_eq!(2, list_array.value_length(2));
875        assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0));
876        assert_eq!(
877            0,
878            unsafe { list_array.value_unchecked(0) }
879                .as_primitive::<Int32Type>()
880                .value(0)
881        );
882        for i in 0..3 {
883            assert!(list_array.is_valid(i));
884            assert!(!list_array.is_null(i));
885        }
886
887        // Now test with a non-zero offset
888        //  [[3, 4, 5], [6, 7]]
889        let list_data = ArrayData::builder(list_data_type)
890            .len(2)
891            .offset(1)
892            .add_buffer(value_offsets)
893            .add_child_data(value_data.clone())
894            .build()
895            .unwrap();
896        let list_array = LargeListArray::from(list_data);
897
898        let values = list_array.values();
899        assert_eq!(value_data, values.to_data());
900        assert_eq!(DataType::Int32, list_array.value_type());
901        assert_eq!(2, list_array.len());
902        assert_eq!(0, list_array.null_count());
903        assert_eq!(6, list_array.value_offsets()[1]);
904        assert_eq!(2, list_array.value_length(1));
905        assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0));
906        assert_eq!(
907            3,
908            unsafe { list_array.value_unchecked(0) }
909                .as_primitive::<Int32Type>()
910                .value(0)
911        );
912    }
913
914    #[test]
915    fn test_list_array_slice() {
916        // Construct a value array
917        let value_data = ArrayData::builder(DataType::Int32)
918            .len(10)
919            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
920            .build()
921            .unwrap();
922
923        // Construct a buffer for value offsets, for the nested array:
924        //  [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
925        let value_offsets = Buffer::from_slice_ref([0, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
926        // 01011001 00000001
927        let mut null_bits: [u8; 2] = [0; 2];
928        bit_util::set_bit(&mut null_bits, 0);
929        bit_util::set_bit(&mut null_bits, 3);
930        bit_util::set_bit(&mut null_bits, 4);
931        bit_util::set_bit(&mut null_bits, 6);
932        bit_util::set_bit(&mut null_bits, 8);
933
934        // Construct a list array from the above two
935        let list_data_type =
936            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
937        let list_data = ArrayData::builder(list_data_type)
938            .len(9)
939            .add_buffer(value_offsets)
940            .add_child_data(value_data.clone())
941            .null_bit_buffer(Some(Buffer::from(null_bits)))
942            .build()
943            .unwrap();
944        let list_array = ListArray::from(list_data);
945
946        let values = list_array.values();
947        assert_eq!(value_data, values.to_data());
948        assert_eq!(DataType::Int32, list_array.value_type());
949        assert_eq!(9, list_array.len());
950        assert_eq!(4, list_array.null_count());
951        assert_eq!(2, list_array.value_offsets()[3]);
952        assert_eq!(2, list_array.value_length(3));
953
954        let sliced_array = list_array.slice(1, 6);
955        assert_eq!(6, sliced_array.len());
956        assert_eq!(3, sliced_array.null_count());
957
958        for i in 0..sliced_array.len() {
959            if bit_util::get_bit(&null_bits, 1 + i) {
960                assert!(sliced_array.is_valid(i));
961            } else {
962                assert!(sliced_array.is_null(i));
963            }
964        }
965
966        // Check offset and length for each non-null value.
967        let sliced_list_array = sliced_array.as_any().downcast_ref::<ListArray>().unwrap();
968        assert_eq!(2, sliced_list_array.value_offsets()[2]);
969        assert_eq!(2, sliced_list_array.value_length(2));
970        assert_eq!(4, sliced_list_array.value_offsets()[3]);
971        assert_eq!(2, sliced_list_array.value_length(3));
972        assert_eq!(6, sliced_list_array.value_offsets()[5]);
973        assert_eq!(3, sliced_list_array.value_length(5));
974    }
975
976    #[test]
977    fn test_large_list_array_slice() {
978        // Construct a value array
979        let value_data = ArrayData::builder(DataType::Int32)
980            .len(10)
981            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
982            .build()
983            .unwrap();
984
985        // Construct a buffer for value offsets, for the nested array:
986        //  [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
987        let value_offsets = Buffer::from_slice_ref([0i64, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
988        // 01011001 00000001
989        let mut null_bits: [u8; 2] = [0; 2];
990        bit_util::set_bit(&mut null_bits, 0);
991        bit_util::set_bit(&mut null_bits, 3);
992        bit_util::set_bit(&mut null_bits, 4);
993        bit_util::set_bit(&mut null_bits, 6);
994        bit_util::set_bit(&mut null_bits, 8);
995
996        // Construct a list array from the above two
997        let list_data_type = DataType::new_large_list(DataType::Int32, false);
998        let list_data = ArrayData::builder(list_data_type)
999            .len(9)
1000            .add_buffer(value_offsets)
1001            .add_child_data(value_data.clone())
1002            .null_bit_buffer(Some(Buffer::from(null_bits)))
1003            .build()
1004            .unwrap();
1005        let list_array = LargeListArray::from(list_data);
1006
1007        let values = list_array.values();
1008        assert_eq!(value_data, values.to_data());
1009        assert_eq!(DataType::Int32, list_array.value_type());
1010        assert_eq!(9, list_array.len());
1011        assert_eq!(4, list_array.null_count());
1012        assert_eq!(2, list_array.value_offsets()[3]);
1013        assert_eq!(2, list_array.value_length(3));
1014
1015        let sliced_array = list_array.slice(1, 6);
1016        assert_eq!(6, sliced_array.len());
1017        assert_eq!(3, sliced_array.null_count());
1018
1019        for i in 0..sliced_array.len() {
1020            if bit_util::get_bit(&null_bits, 1 + i) {
1021                assert!(sliced_array.is_valid(i));
1022            } else {
1023                assert!(sliced_array.is_null(i));
1024            }
1025        }
1026
1027        // Check offset and length for each non-null value.
1028        let sliced_list_array = sliced_array
1029            .as_any()
1030            .downcast_ref::<LargeListArray>()
1031            .unwrap();
1032        assert_eq!(2, sliced_list_array.value_offsets()[2]);
1033        assert_eq!(2, sliced_list_array.value_length(2));
1034        assert_eq!(4, sliced_list_array.value_offsets()[3]);
1035        assert_eq!(2, sliced_list_array.value_length(3));
1036        assert_eq!(6, sliced_list_array.value_offsets()[5]);
1037        assert_eq!(3, sliced_list_array.value_length(5));
1038    }
1039
1040    #[test]
1041    #[should_panic(expected = "index out of bounds: the len is 10 but the index is 11")]
1042    fn test_list_array_index_out_of_bound() {
1043        // Construct a value array
1044        let value_data = ArrayData::builder(DataType::Int32)
1045            .len(10)
1046            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
1047            .build()
1048            .unwrap();
1049
1050        // Construct a buffer for value offsets, for the nested array:
1051        //  [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
1052        let value_offsets = Buffer::from_slice_ref([0i64, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
1053        // 01011001 00000001
1054        let mut null_bits: [u8; 2] = [0; 2];
1055        bit_util::set_bit(&mut null_bits, 0);
1056        bit_util::set_bit(&mut null_bits, 3);
1057        bit_util::set_bit(&mut null_bits, 4);
1058        bit_util::set_bit(&mut null_bits, 6);
1059        bit_util::set_bit(&mut null_bits, 8);
1060
1061        // Construct a list array from the above two
1062        let list_data_type = DataType::new_large_list(DataType::Int32, false);
1063        let list_data = ArrayData::builder(list_data_type)
1064            .len(9)
1065            .add_buffer(value_offsets)
1066            .add_child_data(value_data)
1067            .null_bit_buffer(Some(Buffer::from(null_bits)))
1068            .build()
1069            .unwrap();
1070        let list_array = LargeListArray::from(list_data);
1071        assert_eq!(9, list_array.len());
1072
1073        list_array.value(10);
1074    }
1075    #[test]
1076    #[should_panic(expected = "ListArray data should contain a single buffer only (value offsets)")]
1077    // Different error messages, so skip for now
1078    // https://github.com/apache/arrow-rs/issues/1545
1079    #[cfg(not(feature = "force_validate"))]
1080    fn test_list_array_invalid_buffer_len() {
1081        let value_data = unsafe {
1082            ArrayData::builder(DataType::Int32)
1083                .len(8)
1084                .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
1085                .build_unchecked()
1086        };
1087        let list_data_type =
1088            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1089        let list_data = unsafe {
1090            ArrayData::builder(list_data_type)
1091                .len(3)
1092                .add_child_data(value_data)
1093                .build_unchecked()
1094        };
1095        drop(ListArray::from(list_data));
1096    }
1097
1098    #[test]
1099    #[should_panic(expected = "ListArray should contain a single child array (values array)")]
1100    // Different error messages, so skip for now
1101    // https://github.com/apache/arrow-rs/issues/1545
1102    #[cfg(not(feature = "force_validate"))]
1103    fn test_list_array_invalid_child_array_len() {
1104        let value_offsets = Buffer::from_slice_ref([0, 2, 5, 7]);
1105        let list_data_type =
1106            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1107        let list_data = unsafe {
1108            ArrayData::builder(list_data_type)
1109                .len(3)
1110                .add_buffer(value_offsets)
1111                .build_unchecked()
1112        };
1113        drop(ListArray::from(list_data));
1114    }
1115
1116    #[test]
1117    #[should_panic(expected = "[Large]ListArray's datatype must be [Large]ListArray(). It is List")]
1118    fn test_from_array_data_validation() {
1119        let mut builder = ListBuilder::new(Int32Builder::new());
1120        builder.values().append_value(1);
1121        builder.append(true);
1122        let array = builder.finish();
1123        let _ = LargeListArray::from(array.into_data());
1124    }
1125
1126    #[test]
1127    fn test_list_array_offsets_need_not_start_at_zero() {
1128        let value_data = ArrayData::builder(DataType::Int32)
1129            .len(8)
1130            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
1131            .build()
1132            .unwrap();
1133
1134        let value_offsets = Buffer::from_slice_ref([2, 2, 5, 7]);
1135
1136        let list_data_type =
1137            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1138        let list_data = ArrayData::builder(list_data_type)
1139            .len(3)
1140            .add_buffer(value_offsets)
1141            .add_child_data(value_data)
1142            .build()
1143            .unwrap();
1144
1145        let list_array = ListArray::from(list_data);
1146        assert_eq!(list_array.value_length(0), 0);
1147        assert_eq!(list_array.value_length(1), 3);
1148        assert_eq!(list_array.value_length(2), 2);
1149    }
1150
1151    #[test]
1152    #[should_panic(expected = "Memory pointer is not aligned with the specified scalar type")]
1153    // Different error messages, so skip for now
1154    // https://github.com/apache/arrow-rs/issues/1545
1155    #[cfg(not(feature = "force_validate"))]
1156    fn test_primitive_array_alignment() {
1157        let buf = Buffer::from_slice_ref([0_u64]);
1158        let buf2 = buf.slice(1);
1159        let array_data = unsafe {
1160            ArrayData::builder(DataType::Int32)
1161                .add_buffer(buf2)
1162                .build_unchecked()
1163        };
1164        drop(Int32Array::from(array_data));
1165    }
1166
1167    #[test]
1168    #[should_panic(expected = "Memory pointer is not aligned with the specified scalar type")]
1169    // Different error messages, so skip for now
1170    // https://github.com/apache/arrow-rs/issues/1545
1171    #[cfg(not(feature = "force_validate"))]
1172    fn test_list_array_alignment() {
1173        let buf = Buffer::from_slice_ref([0_u64]);
1174        let buf2 = buf.slice(1);
1175
1176        let values: [i32; 8] = [0; 8];
1177        let value_data = unsafe {
1178            ArrayData::builder(DataType::Int32)
1179                .add_buffer(Buffer::from_slice_ref(values))
1180                .build_unchecked()
1181        };
1182
1183        let list_data_type =
1184            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1185        let list_data = unsafe {
1186            ArrayData::builder(list_data_type)
1187                .add_buffer(buf2)
1188                .add_child_data(value_data)
1189                .build_unchecked()
1190        };
1191        drop(ListArray::from(list_data));
1192    }
1193
1194    #[test]
1195    fn list_array_equality() {
1196        // test scaffold
1197        fn do_comparison(
1198            lhs_data: Vec<Option<Vec<Option<i32>>>>,
1199            rhs_data: Vec<Option<Vec<Option<i32>>>>,
1200            should_equal: bool,
1201        ) {
1202            let lhs = ListArray::from_iter_primitive::<Int32Type, _, _>(lhs_data.clone());
1203            let rhs = ListArray::from_iter_primitive::<Int32Type, _, _>(rhs_data.clone());
1204            assert_eq!(lhs == rhs, should_equal);
1205
1206            let lhs = LargeListArray::from_iter_primitive::<Int32Type, _, _>(lhs_data);
1207            let rhs = LargeListArray::from_iter_primitive::<Int32Type, _, _>(rhs_data);
1208            assert_eq!(lhs == rhs, should_equal);
1209        }
1210
1211        do_comparison(
1212            vec![
1213                Some(vec![Some(0), Some(1), Some(2)]),
1214                None,
1215                Some(vec![Some(3), None, Some(5)]),
1216                Some(vec![Some(6), Some(7)]),
1217            ],
1218            vec![
1219                Some(vec![Some(0), Some(1), Some(2)]),
1220                None,
1221                Some(vec![Some(3), None, Some(5)]),
1222                Some(vec![Some(6), Some(7)]),
1223            ],
1224            true,
1225        );
1226
1227        do_comparison(
1228            vec![
1229                None,
1230                None,
1231                Some(vec![Some(3), None, Some(5)]),
1232                Some(vec![Some(6), Some(7)]),
1233            ],
1234            vec![
1235                Some(vec![Some(0), Some(1), Some(2)]),
1236                None,
1237                Some(vec![Some(3), None, Some(5)]),
1238                Some(vec![Some(6), Some(7)]),
1239            ],
1240            false,
1241        );
1242
1243        do_comparison(
1244            vec![
1245                None,
1246                None,
1247                Some(vec![Some(3), None, Some(5)]),
1248                Some(vec![Some(6), Some(7)]),
1249            ],
1250            vec![
1251                None,
1252                None,
1253                Some(vec![Some(3), None, Some(5)]),
1254                Some(vec![Some(0), Some(0)]),
1255            ],
1256            false,
1257        );
1258
1259        do_comparison(
1260            vec![None, None, Some(vec![Some(1)])],
1261            vec![None, None, Some(vec![Some(2)])],
1262            false,
1263        );
1264    }
1265
1266    #[test]
1267    fn test_empty_offsets() {
1268        let f = Arc::new(Field::new("element", DataType::Int32, true));
1269        let string = ListArray::from(
1270            ArrayData::builder(DataType::List(f.clone()))
1271                .buffers(vec![Buffer::from(&[])])
1272                .add_child_data(ArrayData::new_empty(&DataType::Int32))
1273                .build()
1274                .unwrap(),
1275        );
1276        assert_eq!(string.value_offsets(), &[0]);
1277        let string = LargeListArray::from(
1278            ArrayData::builder(DataType::LargeList(f))
1279                .buffers(vec![Buffer::from(&[])])
1280                .add_child_data(ArrayData::new_empty(&DataType::Int32))
1281                .build()
1282                .unwrap(),
1283        );
1284        assert_eq!(string.len(), 0);
1285        assert_eq!(string.value_offsets(), &[0]);
1286    }
1287
1288    #[test]
1289    fn test_try_new() {
1290        let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
1291        let values = Int32Array::new(vec![1, 2, 3, 4, 5].into(), None);
1292        let values = Arc::new(values) as ArrayRef;
1293
1294        let field = Arc::new(Field::new("element", DataType::Int32, false));
1295        ListArray::new(field.clone(), offsets.clone(), values.clone(), None);
1296
1297        let nulls = NullBuffer::new_null(3);
1298        ListArray::new(field.clone(), offsets, values.clone(), Some(nulls));
1299
1300        let nulls = NullBuffer::new_null(3);
1301        let offsets = OffsetBuffer::new(vec![0, 1, 2, 4, 5].into());
1302        let err = LargeListArray::try_new(field, offsets.clone(), values.clone(), Some(nulls))
1303            .unwrap_err();
1304
1305        assert_eq!(
1306            err.to_string(),
1307            "Invalid argument error: Incorrect length of null buffer for LargeListArray, expected 4 got 3"
1308        );
1309
1310        let field = Arc::new(Field::new("element", DataType::Int64, false));
1311        let err = LargeListArray::try_new(field.clone(), offsets.clone(), values.clone(), None)
1312            .unwrap_err();
1313
1314        assert_eq!(
1315            err.to_string(),
1316            "Invalid argument error: LargeListArray expected data type Int64 got Int32 for \"element\""
1317        );
1318
1319        let nulls = NullBuffer::new_null(7);
1320        let values = Int64Array::new(vec![0; 7].into(), Some(nulls));
1321        let values = Arc::new(values);
1322
1323        let err =
1324            LargeListArray::try_new(field, offsets.clone(), values.clone(), None).unwrap_err();
1325
1326        assert_eq!(
1327            err.to_string(),
1328            "Invalid argument error: Non-nullable field of LargeListArray \"element\" cannot contain nulls"
1329        );
1330
1331        let field = Arc::new(Field::new("element", DataType::Int64, true));
1332        LargeListArray::new(field.clone(), offsets.clone(), values, None);
1333
1334        let values = Int64Array::new(vec![0; 2].into(), None);
1335        let err = LargeListArray::try_new(field, offsets, Arc::new(values), None).unwrap_err();
1336
1337        assert_eq!(
1338            err.to_string(),
1339            "Invalid argument error: Max offset of 5 exceeds length of values 2"
1340        );
1341    }
1342
1343    #[test]
1344    fn test_try_new_non_nullable_field_dictionary_values() {
1345        let keys = Int8Array::new(vec![0i8, 1].into(), Some(NullBuffer::new_valid(2)));
1346        let values = StringArray::from(vec!["x", "y"]);
1347        let dict = Int8DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
1348        let field = Arc::new(Field::new("element", dict.data_type().clone(), false));
1349        ListArray::try_new(
1350            field,
1351            OffsetBuffer::new(vec![0, 2].into()),
1352            Arc::new(dict),
1353            None,
1354        )
1355        .unwrap();
1356
1357        let keys = Int8Array::from(vec![Some(0i8), None]);
1358        let values = StringArray::from(vec!["x", "y"]);
1359        let dict = Int8DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
1360        let field = Arc::new(Field::new("element", dict.data_type().clone(), false));
1361        let err = ListArray::try_new(
1362            field,
1363            OffsetBuffer::new(vec![0, 2].into()),
1364            Arc::new(dict),
1365            None,
1366        )
1367        .unwrap_err();
1368
1369        assert!(err.to_string().contains("cannot contain nulls"));
1370    }
1371
1372    #[test]
1373    fn test_from_fixed_size_list() {
1374        let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
1375        builder.values().append_slice(&[1, 2, 3]);
1376        builder.append(true);
1377        builder.values().append_slice(&[0, 0, 0]);
1378        builder.append(false);
1379        builder.values().append_slice(&[4, 5, 6]);
1380        builder.append(true);
1381        let list: ListArray = builder.finish().into();
1382
1383        let values: Vec<_> = list
1384            .iter()
1385            .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
1386            .collect();
1387        assert_eq!(values, vec![Some(vec![1, 2, 3]), None, Some(vec![4, 5, 6])])
1388    }
1389
1390    #[test]
1391    fn test_nullable_union() {
1392        let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
1393        let mut builder = UnionBuilder::new_dense();
1394        builder.append::<Int32Type>("a", 1).unwrap();
1395        builder.append::<Int32Type>("b", 2).unwrap();
1396        builder.append::<Int32Type>("b", 3).unwrap();
1397        builder.append::<Int32Type>("a", 4).unwrap();
1398        builder.append::<Int32Type>("a", 5).unwrap();
1399        let values = builder.build().unwrap();
1400        let field = Arc::new(Field::new("element", values.data_type().clone(), false));
1401        ListArray::new(field.clone(), offsets, Arc::new(values), None);
1402    }
1403
1404    #[test]
1405    fn test_list_new_null_len() {
1406        let field = Arc::new(Field::new_list_field(DataType::Int32, true));
1407        let array = ListArray::new_null(field, 5);
1408        assert_eq!(array.len(), 5);
1409    }
1410
1411    #[test]
1412    fn test_list_from_iter_i32() {
1413        let array = ListArray::from_nested_iter::<Int32Builder, _, _, _>(vec![
1414            None,
1415            Some(vec![Some(1), None, Some(2)]),
1416        ]);
1417        let expected_offsets = &[0, 0, 3];
1418        let expected_values: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(2)]));
1419        assert_eq!(array.value_offsets(), expected_offsets);
1420        assert_eq!(array.values(), &expected_values);
1421    }
1422
1423    #[test]
1424    fn test_list_from_iter_bool() {
1425        let array = ListArray::from_nested_iter::<BooleanBuilder, _, _, _>(vec![
1426            Some(vec![None, Some(false), Some(true)]),
1427            None,
1428        ]);
1429        let expected_offsets = &[0, 3, 3];
1430        let expected_values: ArrayRef =
1431            Arc::new(BooleanArray::from(vec![None, Some(false), Some(true)]));
1432        assert_eq!(array.value_offsets(), expected_offsets);
1433        assert_eq!(array.values(), &expected_values);
1434    }
1435
1436    #[test]
1437    fn test_list_from_iter_str() {
1438        let array = ListArray::from_nested_iter::<StringBuilder, _, _, _>(vec![
1439            Some(vec![Some("foo"), None, Some("bar")]),
1440            None,
1441        ]);
1442        let expected_offsets = &[0, 3, 3];
1443        let expected_values: ArrayRef =
1444            Arc::new(StringArray::from(vec![Some("foo"), None, Some("bar")]));
1445        assert_eq!(array.value_offsets(), expected_offsets);
1446        assert_eq!(array.values(), &expected_values);
1447    }
1448
1449    #[test]
1450    fn test_list_from_iter_dict_str() {
1451        let array =
1452            ListArray::from_nested_iter::<StringDictionaryBuilder<Int8Type>, _, _, _>(vec![
1453                Some(vec![Some("foo"), None, Some("bar"), Some("foo")]),
1454                None,
1455            ]);
1456        let expected_offsets = &[0, 4, 4];
1457        let expected_dict_values: ArrayRef =
1458            Arc::new(StringArray::from(vec![Some("foo"), Some("bar")]));
1459        let expected_dict_keys = Int8Array::from(vec![Some(0), None, Some(1), Some(0)]);
1460        let expected_values: ArrayRef = Arc::new(
1461            Int8DictionaryArray::try_new(expected_dict_keys, expected_dict_values).unwrap(),
1462        );
1463        assert_eq!(array.value_offsets(), expected_offsets);
1464        assert_eq!(array.values(), &expected_values);
1465    }
1466}
arrow_array/array/list_array.rs

arrow_array/array/
list_array.rs
