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            && n.len() != len
228        {
229            return Err(ArrowError::InvalidArgumentError(format!(
230                "Incorrect length of null buffer for {}ListArray, expected {len} got {}",
231                OffsetSize::PREFIX,
232                n.len(),
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    /// The field that describes the values of this list.
326    pub fn value_field(&self) -> &FieldRef {
327        match &self.data_type {
328            DataType::List(f) | DataType::LargeList(f) => f,
329            _ => unreachable!(),
330        }
331    }
332
333    /// Returns a reference to the offsets of this list
334    ///
335    /// Unlike [`Self::value_offsets`] this returns the [`OffsetBuffer`]
336    /// allowing for zero-copy cloning.
337    ///
338    /// Notes: The `offsets` may not start at 0 and may not cover all values in
339    /// [`Self::values`]. This can happen when the list array was sliced via
340    /// [`Self::slice`]. See documentation for [`Self`] for more details.
341    #[inline]
342    pub fn offsets(&self) -> &OffsetBuffer<OffsetSize> {
343        &self.value_offsets
344    }
345
346    /// Returns a reference to the values of this list
347    ///
348    /// Note: The list array may not refer to all values in the `values` array.
349    /// For example if the list array was sliced via [`Self::slice`] values will
350    /// still contain values both before and after the slice. See documentation
351    /// for [`Self`] for more details.
352    #[inline]
353    pub fn values(&self) -> &ArrayRef {
354        &self.values
355    }
356
357    /// Returns a clone of the value type of this list.
358    pub fn value_type(&self) -> DataType {
359        self.values.data_type().clone()
360    }
361
362    /// Returns ith value of this list array.
363    ///
364    /// Note: This method does not check for nulls and the value is arbitrary
365    /// if [`is_null`](Self::is_null) returns true for the index.
366    ///
367    /// # Safety
368    /// Caller must ensure that the index is within the array bounds
369    pub unsafe fn value_unchecked(&self, i: usize) -> ArrayRef {
370        let end = unsafe { self.value_offsets().get_unchecked(i + 1).as_usize() };
371        let start = unsafe { self.value_offsets().get_unchecked(i).as_usize() };
372        self.values.slice(start, end - start)
373    }
374
375    /// Returns ith value of this list array.
376    ///
377    /// Note: This method does not check for nulls and the value is arbitrary
378    /// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
379    ///
380    /// # Panics
381    /// Panics if index `i` is out of bounds
382    pub fn value(&self, i: usize) -> ArrayRef {
383        let end = self.value_offsets()[i + 1].as_usize();
384        let start = self.value_offsets()[i].as_usize();
385        self.values.slice(start, end - start)
386    }
387
388    /// Returns the offset values in the offsets buffer.
389    ///
390    /// See [`Self::offsets`] for more details.
391    #[inline]
392    pub fn value_offsets(&self) -> &[OffsetSize] {
393        &self.value_offsets
394    }
395
396    /// Returns the length for value at index `i`.
397    ///
398    /// # Panics
399    /// Panics if `i >= self.len()`
400    #[inline]
401    pub fn value_length(&self, i: usize) -> OffsetSize {
402        let offsets = self.value_offsets();
403        offsets[i + 1] - offsets[i]
404    }
405
406    /// constructs a new iterator
407    pub fn iter<'a>(&'a self) -> GenericListArrayIter<'a, OffsetSize> {
408        GenericListArrayIter::<'a, OffsetSize>::new(self)
409    }
410
411    #[inline]
412    fn get_type(data_type: &DataType) -> Option<&DataType> {
413        match (OffsetSize::IS_LARGE, data_type) {
414            (true, DataType::LargeList(child)) | (false, DataType::List(child)) => {
415                Some(child.data_type())
416            }
417            _ => None,
418        }
419    }
420
421    /// Returns a zero-copy slice of this array with the indicated offset and length.
422    ///
423    /// Notes: this method does *NOT* slice the underlying values array or modify
424    /// the values in the offsets buffer. See [`Self::values`] and
425    /// [`Self::offsets`] for more information.
426    ///
427    /// # Panics
428    /// Panics if `offset + length > self.len()`
429    pub fn slice(&self, offset: usize, length: usize) -> Self {
430        Self {
431            data_type: self.data_type.clone(),
432            nulls: self.nulls.as_ref().map(|n| n.slice(offset, length)),
433            values: self.values.clone(),
434            value_offsets: self.value_offsets.slice(offset, length),
435        }
436    }
437
438    /// Creates a [`GenericListArray`] from an iterator of primitive values
439    /// # Example
440    /// ```
441    /// # use arrow_array::ListArray;
442    /// # use arrow_array::types::Int32Type;
443    ///
444    /// let data = vec![
445    ///    Some(vec![Some(0), Some(1), Some(2)]),
446    ///    None,
447    ///    Some(vec![Some(3), None, Some(5)]),
448    ///    Some(vec![Some(6), Some(7)]),
449    /// ];
450    /// let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
451    /// println!("{:?}", list_array);
452    /// ```
453    pub fn from_iter_primitive<T, P, I>(iter: I) -> Self
454    where
455        T: ArrowPrimitiveType,
456        P: IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
457        I: IntoIterator<Item = Option<P>>,
458    {
459        Self::from_nested_iter::<PrimitiveBuilder<T>, T::Native, P, I>(iter)
460    }
461
462    /// Creates a [`GenericListArray`] from a nested iterator of values.
463    /// This method works for any values type that has a corresponding builder that implements the
464    /// `Extend` trait. That includes all numeric types, booleans, binary and string types and also
465    /// dictionary encoded binary and strings.
466    ///
467    /// # Example
468    /// ```
469    /// # use arrow_array::ListArray;
470    /// # use arrow_array::types::Int32Type;
471    /// # use arrow_array::builder::StringDictionaryBuilder;
472    /// let data = vec![
473    ///    Some(vec![Some("foo"), Some("bar"), Some("baz")]),
474    ///    None,
475    ///    Some(vec![Some("bar"), None, Some("foo")]),
476    ///    Some(vec![]),
477    /// ];
478    /// let list_array = ListArray::from_nested_iter::<StringDictionaryBuilder<Int32Type>, _, _, _>(data);
479    /// println!("{:?}", list_array);
480    /// ```
481    pub fn from_nested_iter<B, T, P, I>(iter: I) -> Self
482    where
483        B: ArrayBuilder + Default + Extend<Option<T>>,
484        P: IntoIterator<Item = Option<T>>,
485        I: IntoIterator<Item = Option<P>>,
486    {
487        let iter = iter.into_iter();
488        let size_hint = iter.size_hint().0;
489        let mut builder = GenericListBuilder::with_capacity(B::default(), size_hint);
490
491        for i in iter {
492            match i {
493                Some(p) => {
494                    builder.values().extend(p);
495                    builder.append(true);
496                }
497                None => builder.append(false),
498            }
499        }
500        builder.finish()
501    }
502}
503
504impl<OffsetSize: OffsetSizeTrait> From<ArrayData> for GenericListArray<OffsetSize> {
505    fn from(data: ArrayData) -> Self {
506        Self::try_new_from_array_data(data)
507            .expect("Expected infallible creation of GenericListArray from ArrayDataRef failed")
508    }
509}
510
511impl<OffsetSize: OffsetSizeTrait> From<GenericListArray<OffsetSize>> for ArrayData {
512    fn from(array: GenericListArray<OffsetSize>) -> Self {
513        let len = array.len();
514        let builder = ArrayDataBuilder::new(array.data_type)
515            .len(len)
516            .nulls(array.nulls)
517            .buffers(vec![array.value_offsets.into_inner().into_inner()])
518            .child_data(vec![array.values.to_data()]);
519
520        unsafe { builder.build_unchecked() }
521    }
522}
523
524impl<OffsetSize: OffsetSizeTrait> From<FixedSizeListArray> for GenericListArray<OffsetSize> {
525    fn from(value: FixedSizeListArray) -> Self {
526        let (field, size) = match value.data_type() {
527            DataType::FixedSizeList(f, size) => (f, *size as usize),
528            _ => unreachable!(),
529        };
530
531        let offsets = OffsetBuffer::from_repeated_length(size, value.len());
532
533        Self {
534            data_type: Self::DATA_TYPE_CONSTRUCTOR(field.clone()),
535            nulls: value.nulls().cloned(),
536            values: value.values().clone(),
537            value_offsets: offsets,
538        }
539    }
540}
541
542impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
543    fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
544        let (data_type, len, nulls, offset, mut buffers, mut child_data) = data.into_parts();
545
546        if buffers.len() != 1 {
547            return Err(ArrowError::InvalidArgumentError(format!(
548                "ListArray data should contain a single buffer only (value offsets), had {}",
549                buffers.len()
550            )));
551        }
552        let buffer = buffers.pop().expect("checked above");
553
554        if child_data.len() != 1 {
555            return Err(ArrowError::InvalidArgumentError(format!(
556                "ListArray should contain a single child array (values array), had {}",
557                child_data.len()
558            )));
559        }
560
561        let values = child_data.pop().expect("checked above");
562
563        if let Some(child_data_type) = Self::get_type(&data_type) {
564            if values.data_type() != child_data_type {
565                return Err(ArrowError::InvalidArgumentError(format!(
566                    "[Large]ListArray's child datatype {:?} does not \
567                             correspond to the List's datatype {:?}",
568                    values.data_type(),
569                    child_data_type
570                )));
571            }
572        } else {
573            return Err(ArrowError::InvalidArgumentError(format!(
574                "[Large]ListArray's datatype must be [Large]ListArray(). It is {data_type:?}",
575            )));
576        }
577
578        let values = make_array(values);
579        // SAFETY:
580        // ArrayData is valid, and verified type above
581        let value_offsets = unsafe { get_offsets_from_buffer(buffer, offset, len) };
582
583        Ok(Self {
584            data_type,
585            nulls,
586            values,
587            value_offsets,
588        })
589    }
590}
591
592/// SAFETY: Correctly implements the contract of Arrow Arrays
593unsafe impl<OffsetSize: OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
594    fn as_any(&self) -> &dyn Any {
595        self
596    }
597
598    fn to_data(&self) -> ArrayData {
599        self.clone().into()
600    }
601
602    fn into_data(self) -> ArrayData {
603        self.into()
604    }
605
606    fn data_type(&self) -> &DataType {
607        &self.data_type
608    }
609
610    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
611        Arc::new(self.slice(offset, length))
612    }
613
614    fn len(&self) -> usize {
615        self.value_offsets.len() - 1
616    }
617
618    fn is_empty(&self) -> bool {
619        self.value_offsets.len() <= 1
620    }
621
622    fn shrink_to_fit(&mut self) {
623        if let Some(nulls) = &mut self.nulls {
624            nulls.shrink_to_fit();
625        }
626        self.values.shrink_to_fit();
627        self.value_offsets.shrink_to_fit();
628    }
629
630    fn offset(&self) -> usize {
631        0
632    }
633
634    fn nulls(&self) -> Option<&NullBuffer> {
635        self.nulls.as_ref()
636    }
637
638    fn logical_null_count(&self) -> usize {
639        // More efficient that the default implementation
640        self.null_count()
641    }
642
643    fn get_buffer_memory_size(&self) -> usize {
644        let mut size = self.values.get_buffer_memory_size();
645        size += self.value_offsets.inner().inner().capacity();
646        if let Some(n) = self.nulls.as_ref() {
647            size += n.buffer().capacity();
648        }
649        size
650    }
651
652    fn get_array_memory_size(&self) -> usize {
653        let mut size = std::mem::size_of::<Self>() + self.values.get_array_memory_size();
654        size += self.value_offsets.inner().inner().capacity();
655        if let Some(n) = self.nulls.as_ref() {
656            size += n.buffer().capacity();
657        }
658        size
659    }
660
661    #[cfg(feature = "pool")]
662    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
663        self.value_offsets.claim(pool);
664        self.values.claim(pool);
665        if let Some(nulls) = &self.nulls {
666            nulls.claim(pool);
667        }
668    }
669}
670
671impl<OffsetSize: OffsetSizeTrait> super::ListLikeArray for GenericListArray<OffsetSize> {
672    fn values(&self) -> &ArrayRef {
673        self.values()
674    }
675
676    fn element_range(&self, index: usize) -> std::ops::Range<usize> {
677        let offsets = self.offsets();
678        let start = offsets[index].as_usize();
679        let end = offsets[index + 1].as_usize();
680        start..end
681    }
682}
683
684impl<OffsetSize: OffsetSizeTrait> ArrayAccessor for &GenericListArray<OffsetSize> {
685    type Item = ArrayRef;
686
687    fn value(&self, index: usize) -> Self::Item {
688        GenericListArray::value(self, index)
689    }
690
691    unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
692        GenericListArray::value(self, index)
693    }
694}
695
696impl<OffsetSize: OffsetSizeTrait> std::fmt::Debug for GenericListArray<OffsetSize> {
697    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
698        let prefix = OffsetSize::PREFIX;
699
700        write!(f, "{prefix}ListArray\n[\n")?;
701        print_long_array(self, f, |array, index, f| {
702            std::fmt::Debug::fmt(&array.value(index), f)
703        })?;
704        write!(f, "]")
705    }
706}
707
708/// A [`GenericListArray`] of variable size lists, storing offsets as `i32`.
709///
710/// See [`ListBuilder`](crate::builder::ListBuilder) for how to construct a [`ListArray`]
711pub type ListArray = GenericListArray<i32>;
712
713/// A [`GenericListArray`] of variable size lists, storing offsets as `i64`.
714///
715/// See [`LargeListBuilder`](crate::builder::LargeListBuilder) for how to construct a [`LargeListArray`]
716pub type LargeListArray = GenericListArray<i64>;
717
718#[cfg(test)]
719mod tests {
720    use super::*;
721    use crate::builder::{
722        BooleanBuilder, FixedSizeListBuilder, Int32Builder, ListBuilder, StringBuilder,
723        StringDictionaryBuilder, UnionBuilder,
724    };
725    use crate::cast::AsArray;
726    use crate::types::{Int8Type, Int32Type};
727    use crate::{
728        BooleanArray, Int8Array, Int8DictionaryArray, Int32Array, Int64Array, StringArray,
729    };
730    use arrow_buffer::{Buffer, ScalarBuffer, bit_util};
731    use arrow_schema::Field;
732
733    fn create_from_buffers() -> ListArray {
734        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
735        let values = Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7]);
736        let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 6, 8]));
737        let field = Arc::new(Field::new_list_field(DataType::Int32, true));
738        ListArray::new(field, offsets, Arc::new(values), None)
739    }
740
741    #[test]
742    fn test_from_iter_primitive() {
743        let data = vec![
744            Some(vec![Some(0), Some(1), Some(2)]),
745            Some(vec![Some(3), Some(4), Some(5)]),
746            Some(vec![Some(6), Some(7)]),
747        ];
748        let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
749
750        let another = create_from_buffers();
751        assert_eq!(list_array, another)
752    }
753
754    #[test]
755    fn test_empty_list_array() {
756        // Construct an empty value array
757        let value_data = ArrayData::builder(DataType::Int32)
758            .len(0)
759            .add_buffer(Buffer::from([]))
760            .build()
761            .unwrap();
762
763        // Construct an empty offset buffer
764        let value_offsets = Buffer::from([]);
765
766        // Construct a list array from the above two
767        let list_data_type =
768            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
769        let list_data = ArrayData::builder(list_data_type)
770            .len(0)
771            .add_buffer(value_offsets)
772            .add_child_data(value_data)
773            .build()
774            .unwrap();
775
776        let list_array = ListArray::from(list_data);
777        assert_eq!(list_array.len(), 0)
778    }
779
780    #[test]
781    fn test_list_array() {
782        // Construct a value array
783        let value_data = ArrayData::builder(DataType::Int32)
784            .len(8)
785            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
786            .build()
787            .unwrap();
788
789        // Construct a buffer for value offsets, for the nested array:
790        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
791        let value_offsets = Buffer::from_slice_ref([0, 3, 6, 8]);
792
793        // Construct a list array from the above two
794        let list_data_type =
795            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
796        let list_data = ArrayData::builder(list_data_type.clone())
797            .len(3)
798            .add_buffer(value_offsets.clone())
799            .add_child_data(value_data.clone())
800            .build()
801            .unwrap();
802        let list_array = ListArray::from(list_data);
803
804        let values = list_array.values();
805        assert_eq!(value_data, values.to_data());
806        assert_eq!(DataType::Int32, list_array.value_type());
807        assert_eq!(3, list_array.len());
808        assert_eq!(0, list_array.null_count());
809        assert_eq!(6, list_array.value_offsets()[2]);
810        assert_eq!(2, list_array.value_length(2));
811        assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0));
812        assert_eq!(
813            0,
814            unsafe { list_array.value_unchecked(0) }
815                .as_primitive::<Int32Type>()
816                .value(0)
817        );
818        for i in 0..3 {
819            assert!(list_array.is_valid(i));
820            assert!(!list_array.is_null(i));
821        }
822
823        // Now test with a non-zero offset (skip first element)
824        //  [[3, 4, 5], [6, 7]]
825        let list_data = ArrayData::builder(list_data_type)
826            .len(2)
827            .offset(1)
828            .add_buffer(value_offsets)
829            .add_child_data(value_data.clone())
830            .build()
831            .unwrap();
832        let list_array = ListArray::from(list_data);
833
834        let values = list_array.values();
835        assert_eq!(value_data, values.to_data());
836        assert_eq!(DataType::Int32, list_array.value_type());
837        assert_eq!(2, list_array.len());
838        assert_eq!(0, list_array.null_count());
839        assert_eq!(6, list_array.value_offsets()[1]);
840        assert_eq!(2, list_array.value_length(1));
841        assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0));
842        assert_eq!(
843            3,
844            unsafe { list_array.value_unchecked(0) }
845                .as_primitive::<Int32Type>()
846                .value(0)
847        );
848    }
849
850    #[test]
851    fn test_large_list_array() {
852        // Construct a value array
853        let value_data = ArrayData::builder(DataType::Int32)
854            .len(8)
855            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
856            .build()
857            .unwrap();
858
859        // Construct a buffer for value offsets, for the nested array:
860        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
861        let value_offsets = Buffer::from_slice_ref([0i64, 3, 6, 8]);
862
863        // Construct a list array from the above two
864        let list_data_type = DataType::new_large_list(DataType::Int32, false);
865        let list_data = ArrayData::builder(list_data_type.clone())
866            .len(3)
867            .add_buffer(value_offsets.clone())
868            .add_child_data(value_data.clone())
869            .build()
870            .unwrap();
871        let list_array = LargeListArray::from(list_data);
872
873        let values = list_array.values();
874        assert_eq!(value_data, values.to_data());
875        assert_eq!(DataType::Int32, list_array.value_type());
876        assert_eq!(3, list_array.len());
877        assert_eq!(0, list_array.null_count());
878        assert_eq!(6, list_array.value_offsets()[2]);
879        assert_eq!(2, list_array.value_length(2));
880        assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0));
881        assert_eq!(
882            0,
883            unsafe { list_array.value_unchecked(0) }
884                .as_primitive::<Int32Type>()
885                .value(0)
886        );
887        for i in 0..3 {
888            assert!(list_array.is_valid(i));
889            assert!(!list_array.is_null(i));
890        }
891
892        // Now test with a non-zero offset
893        //  [[3, 4, 5], [6, 7]]
894        let list_data = ArrayData::builder(list_data_type)
895            .len(2)
896            .offset(1)
897            .add_buffer(value_offsets)
898            .add_child_data(value_data.clone())
899            .build()
900            .unwrap();
901        let list_array = LargeListArray::from(list_data);
902
903        let values = list_array.values();
904        assert_eq!(value_data, values.to_data());
905        assert_eq!(DataType::Int32, list_array.value_type());
906        assert_eq!(2, list_array.len());
907        assert_eq!(0, list_array.null_count());
908        assert_eq!(6, list_array.value_offsets()[1]);
909        assert_eq!(2, list_array.value_length(1));
910        assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0));
911        assert_eq!(
912            3,
913            unsafe { list_array.value_unchecked(0) }
914                .as_primitive::<Int32Type>()
915                .value(0)
916        );
917    }
918
919    #[test]
920    fn test_list_array_slice() {
921        // Construct a value array
922        let value_data = ArrayData::builder(DataType::Int32)
923            .len(10)
924            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
925            .build()
926            .unwrap();
927
928        // Construct a buffer for value offsets, for the nested array:
929        //  [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
930        let value_offsets = Buffer::from_slice_ref([0, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
931        // 01011001 00000001
932        let mut null_bits: [u8; 2] = [0; 2];
933        bit_util::set_bit(&mut null_bits, 0);
934        bit_util::set_bit(&mut null_bits, 3);
935        bit_util::set_bit(&mut null_bits, 4);
936        bit_util::set_bit(&mut null_bits, 6);
937        bit_util::set_bit(&mut null_bits, 8);
938
939        // Construct a list array from the above two
940        let list_data_type =
941            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
942        let list_data = ArrayData::builder(list_data_type)
943            .len(9)
944            .add_buffer(value_offsets)
945            .add_child_data(value_data.clone())
946            .null_bit_buffer(Some(Buffer::from(null_bits)))
947            .build()
948            .unwrap();
949        let list_array = ListArray::from(list_data);
950
951        let values = list_array.values();
952        assert_eq!(value_data, values.to_data());
953        assert_eq!(DataType::Int32, list_array.value_type());
954        assert_eq!(9, list_array.len());
955        assert_eq!(4, list_array.null_count());
956        assert_eq!(2, list_array.value_offsets()[3]);
957        assert_eq!(2, list_array.value_length(3));
958
959        let sliced_array = list_array.slice(1, 6);
960        assert_eq!(6, sliced_array.len());
961        assert_eq!(3, sliced_array.null_count());
962
963        for i in 0..sliced_array.len() {
964            if bit_util::get_bit(&null_bits, 1 + i) {
965                assert!(sliced_array.is_valid(i));
966            } else {
967                assert!(sliced_array.is_null(i));
968            }
969        }
970
971        // Check offset and length for each non-null value.
972        let sliced_list_array = sliced_array.as_any().downcast_ref::<ListArray>().unwrap();
973        assert_eq!(2, sliced_list_array.value_offsets()[2]);
974        assert_eq!(2, sliced_list_array.value_length(2));
975        assert_eq!(4, sliced_list_array.value_offsets()[3]);
976        assert_eq!(2, sliced_list_array.value_length(3));
977        assert_eq!(6, sliced_list_array.value_offsets()[5]);
978        assert_eq!(3, sliced_list_array.value_length(5));
979    }
980
981    #[test]
982    fn test_large_list_array_slice() {
983        // Construct a value array
984        let value_data = ArrayData::builder(DataType::Int32)
985            .len(10)
986            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
987            .build()
988            .unwrap();
989
990        // Construct a buffer for value offsets, for the nested array:
991        //  [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
992        let value_offsets = Buffer::from_slice_ref([0i64, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
993        // 01011001 00000001
994        let mut null_bits: [u8; 2] = [0; 2];
995        bit_util::set_bit(&mut null_bits, 0);
996        bit_util::set_bit(&mut null_bits, 3);
997        bit_util::set_bit(&mut null_bits, 4);
998        bit_util::set_bit(&mut null_bits, 6);
999        bit_util::set_bit(&mut null_bits, 8);
1000
1001        // Construct a list array from the above two
1002        let list_data_type = DataType::new_large_list(DataType::Int32, false);
1003        let list_data = ArrayData::builder(list_data_type)
1004            .len(9)
1005            .add_buffer(value_offsets)
1006            .add_child_data(value_data.clone())
1007            .null_bit_buffer(Some(Buffer::from(null_bits)))
1008            .build()
1009            .unwrap();
1010        let list_array = LargeListArray::from(list_data);
1011
1012        let values = list_array.values();
1013        assert_eq!(value_data, values.to_data());
1014        assert_eq!(DataType::Int32, list_array.value_type());
1015        assert_eq!(9, list_array.len());
1016        assert_eq!(4, list_array.null_count());
1017        assert_eq!(2, list_array.value_offsets()[3]);
1018        assert_eq!(2, list_array.value_length(3));
1019
1020        let sliced_array = list_array.slice(1, 6);
1021        assert_eq!(6, sliced_array.len());
1022        assert_eq!(3, sliced_array.null_count());
1023
1024        for i in 0..sliced_array.len() {
1025            if bit_util::get_bit(&null_bits, 1 + i) {
1026                assert!(sliced_array.is_valid(i));
1027            } else {
1028                assert!(sliced_array.is_null(i));
1029            }
1030        }
1031
1032        // Check offset and length for each non-null value.
1033        let sliced_list_array = sliced_array
1034            .as_any()
1035            .downcast_ref::<LargeListArray>()
1036            .unwrap();
1037        assert_eq!(2, sliced_list_array.value_offsets()[2]);
1038        assert_eq!(2, sliced_list_array.value_length(2));
1039        assert_eq!(4, sliced_list_array.value_offsets()[3]);
1040        assert_eq!(2, sliced_list_array.value_length(3));
1041        assert_eq!(6, sliced_list_array.value_offsets()[5]);
1042        assert_eq!(3, sliced_list_array.value_length(5));
1043    }
1044
1045    #[test]
1046    #[should_panic(expected = "index out of bounds: the len is 10 but the index is 11")]
1047    fn test_list_array_index_out_of_bound() {
1048        // Construct a value array
1049        let value_data = ArrayData::builder(DataType::Int32)
1050            .len(10)
1051            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
1052            .build()
1053            .unwrap();
1054
1055        // Construct a buffer for value offsets, for the nested array:
1056        //  [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
1057        let value_offsets = Buffer::from_slice_ref([0i64, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
1058        // 01011001 00000001
1059        let mut null_bits: [u8; 2] = [0; 2];
1060        bit_util::set_bit(&mut null_bits, 0);
1061        bit_util::set_bit(&mut null_bits, 3);
1062        bit_util::set_bit(&mut null_bits, 4);
1063        bit_util::set_bit(&mut null_bits, 6);
1064        bit_util::set_bit(&mut null_bits, 8);
1065
1066        // Construct a list array from the above two
1067        let list_data_type = DataType::new_large_list(DataType::Int32, false);
1068        let list_data = ArrayData::builder(list_data_type)
1069            .len(9)
1070            .add_buffer(value_offsets)
1071            .add_child_data(value_data)
1072            .null_bit_buffer(Some(Buffer::from(null_bits)))
1073            .build()
1074            .unwrap();
1075        let list_array = LargeListArray::from(list_data);
1076        assert_eq!(9, list_array.len());
1077
1078        list_array.value(10);
1079    }
1080    #[test]
1081    #[should_panic(expected = "ListArray data should contain a single buffer only (value offsets)")]
1082    // Different error messages, so skip for now
1083    // https://github.com/apache/arrow-rs/issues/1545
1084    #[cfg(not(feature = "force_validate"))]
1085    fn test_list_array_invalid_buffer_len() {
1086        let value_data = unsafe {
1087            ArrayData::builder(DataType::Int32)
1088                .len(8)
1089                .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
1090                .build_unchecked()
1091        };
1092        let list_data_type =
1093            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1094        let list_data = unsafe {
1095            ArrayData::builder(list_data_type)
1096                .len(3)
1097                .add_child_data(value_data)
1098                .build_unchecked()
1099        };
1100        drop(ListArray::from(list_data));
1101    }
1102
1103    #[test]
1104    #[should_panic(expected = "ListArray should contain a single child array (values array)")]
1105    // Different error messages, so skip for now
1106    // https://github.com/apache/arrow-rs/issues/1545
1107    #[cfg(not(feature = "force_validate"))]
1108    fn test_list_array_invalid_child_array_len() {
1109        let value_offsets = Buffer::from_slice_ref([0, 2, 5, 7]);
1110        let list_data_type =
1111            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1112        let list_data = unsafe {
1113            ArrayData::builder(list_data_type)
1114                .len(3)
1115                .add_buffer(value_offsets)
1116                .build_unchecked()
1117        };
1118        drop(ListArray::from(list_data));
1119    }
1120
1121    #[test]
1122    #[should_panic(expected = "[Large]ListArray's datatype must be [Large]ListArray(). It is List")]
1123    fn test_from_array_data_validation() {
1124        let mut builder = ListBuilder::new(Int32Builder::new());
1125        builder.values().append_value(1);
1126        builder.append(true);
1127        let array = builder.finish();
1128        let _ = LargeListArray::from(array.into_data());
1129    }
1130
1131    #[test]
1132    fn test_list_array_offsets_need_not_start_at_zero() {
1133        let value_data = ArrayData::builder(DataType::Int32)
1134            .len(8)
1135            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
1136            .build()
1137            .unwrap();
1138
1139        let value_offsets = Buffer::from_slice_ref([2, 2, 5, 7]);
1140
1141        let list_data_type =
1142            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1143        let list_data = ArrayData::builder(list_data_type)
1144            .len(3)
1145            .add_buffer(value_offsets)
1146            .add_child_data(value_data)
1147            .build()
1148            .unwrap();
1149
1150        let list_array = ListArray::from(list_data);
1151        assert_eq!(list_array.value_length(0), 0);
1152        assert_eq!(list_array.value_length(1), 3);
1153        assert_eq!(list_array.value_length(2), 2);
1154    }
1155
1156    #[test]
1157    #[should_panic(expected = "Memory pointer is not aligned with the specified scalar type")]
1158    // Different error messages, so skip for now
1159    // https://github.com/apache/arrow-rs/issues/1545
1160    #[cfg(not(feature = "force_validate"))]
1161    fn test_primitive_array_alignment() {
1162        let buf = Buffer::from_slice_ref([0_u64]);
1163        let buf2 = buf.slice(1);
1164        let array_data = unsafe {
1165            ArrayData::builder(DataType::Int32)
1166                .add_buffer(buf2)
1167                .build_unchecked()
1168        };
1169        drop(Int32Array::from(array_data));
1170    }
1171
1172    #[test]
1173    #[should_panic(expected = "Memory pointer is not aligned with the specified scalar type")]
1174    // Different error messages, so skip for now
1175    // https://github.com/apache/arrow-rs/issues/1545
1176    #[cfg(not(feature = "force_validate"))]
1177    fn test_list_array_alignment() {
1178        let buf = Buffer::from_slice_ref([0_u64]);
1179        let buf2 = buf.slice(1);
1180
1181        let values: [i32; 8] = [0; 8];
1182        let value_data = unsafe {
1183            ArrayData::builder(DataType::Int32)
1184                .add_buffer(Buffer::from_slice_ref(values))
1185                .build_unchecked()
1186        };
1187
1188        let list_data_type =
1189            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1190        let list_data = unsafe {
1191            ArrayData::builder(list_data_type)
1192                .add_buffer(buf2)
1193                .add_child_data(value_data)
1194                .build_unchecked()
1195        };
1196        drop(ListArray::from(list_data));
1197    }
1198
1199    #[test]
1200    fn list_array_equality() {
1201        // test scaffold
1202        fn do_comparison(
1203            lhs_data: Vec<Option<Vec<Option<i32>>>>,
1204            rhs_data: Vec<Option<Vec<Option<i32>>>>,
1205            should_equal: bool,
1206        ) {
1207            let lhs = ListArray::from_iter_primitive::<Int32Type, _, _>(lhs_data.clone());
1208            let rhs = ListArray::from_iter_primitive::<Int32Type, _, _>(rhs_data.clone());
1209            assert_eq!(lhs == rhs, should_equal);
1210
1211            let lhs = LargeListArray::from_iter_primitive::<Int32Type, _, _>(lhs_data);
1212            let rhs = LargeListArray::from_iter_primitive::<Int32Type, _, _>(rhs_data);
1213            assert_eq!(lhs == rhs, should_equal);
1214        }
1215
1216        do_comparison(
1217            vec![
1218                Some(vec![Some(0), Some(1), Some(2)]),
1219                None,
1220                Some(vec![Some(3), None, Some(5)]),
1221                Some(vec![Some(6), Some(7)]),
1222            ],
1223            vec![
1224                Some(vec![Some(0), Some(1), Some(2)]),
1225                None,
1226                Some(vec![Some(3), None, Some(5)]),
1227                Some(vec![Some(6), Some(7)]),
1228            ],
1229            true,
1230        );
1231
1232        do_comparison(
1233            vec![
1234                None,
1235                None,
1236                Some(vec![Some(3), None, Some(5)]),
1237                Some(vec![Some(6), Some(7)]),
1238            ],
1239            vec![
1240                Some(vec![Some(0), Some(1), Some(2)]),
1241                None,
1242                Some(vec![Some(3), None, Some(5)]),
1243                Some(vec![Some(6), Some(7)]),
1244            ],
1245            false,
1246        );
1247
1248        do_comparison(
1249            vec![
1250                None,
1251                None,
1252                Some(vec![Some(3), None, Some(5)]),
1253                Some(vec![Some(6), Some(7)]),
1254            ],
1255            vec![
1256                None,
1257                None,
1258                Some(vec![Some(3), None, Some(5)]),
1259                Some(vec![Some(0), Some(0)]),
1260            ],
1261            false,
1262        );
1263
1264        do_comparison(
1265            vec![None, None, Some(vec![Some(1)])],
1266            vec![None, None, Some(vec![Some(2)])],
1267            false,
1268        );
1269    }
1270
1271    #[test]
1272    fn test_empty_offsets() {
1273        let f = Arc::new(Field::new("element", DataType::Int32, true));
1274        let string = ListArray::from(
1275            ArrayData::builder(DataType::List(f.clone()))
1276                .buffers(vec![Buffer::from(&[])])
1277                .add_child_data(ArrayData::new_empty(&DataType::Int32))
1278                .build()
1279                .unwrap(),
1280        );
1281        assert_eq!(string.value_offsets(), &[0]);
1282        let string = LargeListArray::from(
1283            ArrayData::builder(DataType::LargeList(f))
1284                .buffers(vec![Buffer::from(&[])])
1285                .add_child_data(ArrayData::new_empty(&DataType::Int32))
1286                .build()
1287                .unwrap(),
1288        );
1289        assert_eq!(string.len(), 0);
1290        assert_eq!(string.value_offsets(), &[0]);
1291    }
1292
1293    #[test]
1294    fn test_try_new() {
1295        let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
1296        let values = Int32Array::new(vec![1, 2, 3, 4, 5].into(), None);
1297        let values = Arc::new(values) as ArrayRef;
1298
1299        let field = Arc::new(Field::new("element", DataType::Int32, false));
1300        ListArray::new(field.clone(), offsets.clone(), values.clone(), None);
1301
1302        let nulls = NullBuffer::new_null(3);
1303        ListArray::new(field.clone(), offsets, values.clone(), Some(nulls));
1304
1305        let nulls = NullBuffer::new_null(3);
1306        let offsets = OffsetBuffer::new(vec![0, 1, 2, 4, 5].into());
1307        let err = LargeListArray::try_new(field, offsets.clone(), values.clone(), Some(nulls))
1308            .unwrap_err();
1309
1310        assert_eq!(
1311            err.to_string(),
1312            "Invalid argument error: Incorrect length of null buffer for LargeListArray, expected 4 got 3"
1313        );
1314
1315        let field = Arc::new(Field::new("element", DataType::Int64, false));
1316        let err = LargeListArray::try_new(field.clone(), offsets.clone(), values.clone(), None)
1317            .unwrap_err();
1318
1319        assert_eq!(
1320            err.to_string(),
1321            "Invalid argument error: LargeListArray expected data type Int64 got Int32 for \"element\""
1322        );
1323
1324        let nulls = NullBuffer::new_null(7);
1325        let values = Int64Array::new(vec![0; 7].into(), Some(nulls));
1326        let values = Arc::new(values);
1327
1328        let err =
1329            LargeListArray::try_new(field, offsets.clone(), values.clone(), None).unwrap_err();
1330
1331        assert_eq!(
1332            err.to_string(),
1333            "Invalid argument error: Non-nullable field of LargeListArray \"element\" cannot contain nulls"
1334        );
1335
1336        let field = Arc::new(Field::new("element", DataType::Int64, true));
1337        LargeListArray::new(field.clone(), offsets.clone(), values, None);
1338
1339        let values = Int64Array::new(vec![0; 2].into(), None);
1340        let err = LargeListArray::try_new(field, offsets, Arc::new(values), None).unwrap_err();
1341
1342        assert_eq!(
1343            err.to_string(),
1344            "Invalid argument error: Max offset of 5 exceeds length of values 2"
1345        );
1346    }
1347
1348    #[test]
1349    fn test_from_fixed_size_list() {
1350        let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
1351        builder.values().append_slice(&[1, 2, 3]);
1352        builder.append(true);
1353        builder.values().append_slice(&[0, 0, 0]);
1354        builder.append(false);
1355        builder.values().append_slice(&[4, 5, 6]);
1356        builder.append(true);
1357        let list: ListArray = builder.finish().into();
1358
1359        let values: Vec<_> = list
1360            .iter()
1361            .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
1362            .collect();
1363        assert_eq!(values, vec![Some(vec![1, 2, 3]), None, Some(vec![4, 5, 6])])
1364    }
1365
1366    #[test]
1367    fn test_nullable_union() {
1368        let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
1369        let mut builder = UnionBuilder::new_dense();
1370        builder.append::<Int32Type>("a", 1).unwrap();
1371        builder.append::<Int32Type>("b", 2).unwrap();
1372        builder.append::<Int32Type>("b", 3).unwrap();
1373        builder.append::<Int32Type>("a", 4).unwrap();
1374        builder.append::<Int32Type>("a", 5).unwrap();
1375        let values = builder.build().unwrap();
1376        let field = Arc::new(Field::new("element", values.data_type().clone(), false));
1377        ListArray::new(field.clone(), offsets, Arc::new(values), None);
1378    }
1379
1380    #[test]
1381    fn test_list_new_null_len() {
1382        let field = Arc::new(Field::new_list_field(DataType::Int32, true));
1383        let array = ListArray::new_null(field, 5);
1384        assert_eq!(array.len(), 5);
1385    }
1386
1387    #[test]
1388    fn test_list_from_iter_i32() {
1389        let array = ListArray::from_nested_iter::<Int32Builder, _, _, _>(vec![
1390            None,
1391            Some(vec![Some(1), None, Some(2)]),
1392        ]);
1393        let expected_offsets = &[0, 0, 3];
1394        let expected_values: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(2)]));
1395        assert_eq!(array.value_offsets(), expected_offsets);
1396        assert_eq!(array.values(), &expected_values);
1397    }
1398
1399    #[test]
1400    fn test_list_from_iter_bool() {
1401        let array = ListArray::from_nested_iter::<BooleanBuilder, _, _, _>(vec![
1402            Some(vec![None, Some(false), Some(true)]),
1403            None,
1404        ]);
1405        let expected_offsets = &[0, 3, 3];
1406        let expected_values: ArrayRef =
1407            Arc::new(BooleanArray::from(vec![None, Some(false), Some(true)]));
1408        assert_eq!(array.value_offsets(), expected_offsets);
1409        assert_eq!(array.values(), &expected_values);
1410    }
1411
1412    #[test]
1413    fn test_list_from_iter_str() {
1414        let array = ListArray::from_nested_iter::<StringBuilder, _, _, _>(vec![
1415            Some(vec![Some("foo"), None, Some("bar")]),
1416            None,
1417        ]);
1418        let expected_offsets = &[0, 3, 3];
1419        let expected_values: ArrayRef =
1420            Arc::new(StringArray::from(vec![Some("foo"), None, Some("bar")]));
1421        assert_eq!(array.value_offsets(), expected_offsets);
1422        assert_eq!(array.values(), &expected_values);
1423    }
1424
1425    #[test]
1426    fn test_list_from_iter_dict_str() {
1427        let array =
1428            ListArray::from_nested_iter::<StringDictionaryBuilder<Int8Type>, _, _, _>(vec![
1429                Some(vec![Some("foo"), None, Some("bar"), Some("foo")]),
1430                None,
1431            ]);
1432        let expected_offsets = &[0, 4, 4];
1433        let expected_dict_values: ArrayRef =
1434            Arc::new(StringArray::from(vec![Some("foo"), Some("bar")]));
1435        let expected_dict_keys = Int8Array::from(vec![Some(0), None, Some(1), Some(0)]);
1436        let expected_values: ArrayRef = Arc::new(
1437            Int8DictionaryArray::try_new(expected_dict_keys, expected_dict_values).unwrap(),
1438        );
1439        assert_eq!(array.value_offsets(), expected_offsets);
1440        assert_eq!(array.values(), &expected_values);
1441    }
1442}
arrow_array/array/list_array.rs

arrow_array/array/
list_array.rs
