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

arrow_array/array/
list_array.rs
