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

arrow_array/array/
list_array.rs
