arrow_array/array/
fixed_size_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::print_long_array;
19use crate::builder::{FixedSizeListBuilder, PrimitiveBuilder};
20use crate::iterator::FixedSizeListIter;
21use crate::{Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, make_array};
22use arrow_buffer::ArrowNativeType;
23use arrow_buffer::buffer::NullBuffer;
24use arrow_data::{ArrayData, ArrayDataBuilder};
25use arrow_schema::{ArrowError, DataType, FieldRef};
26use std::any::Any;
27use std::sync::Arc;
28
29/// An array of [fixed length lists], similar to JSON arrays
30/// (e.g. `["A", "B"]`).
31///
32/// Lists are represented using a `values` child
33/// array where each list has a fixed size of `value_length`.
34///
35/// Use [`FixedSizeListBuilder`] to construct a [`FixedSizeListArray`].
36///
37/// # Representation
38///
39/// A [`FixedSizeListArray`] can represent a list of values of any other
40/// supported Arrow type. Each element of the `FixedSizeListArray` itself is
41/// a list which may contain NULL and non-null values,
42/// or may itself be NULL.
43///
44/// For example, this `FixedSizeListArray` stores lists of strings:
45///
46/// ```text
47/// ┌─────────────┐
48/// │    [A,B]    │
49/// ├─────────────┤
50/// │    NULL     │
51/// ├─────────────┤
52/// │   [C,NULL]  │
53/// └─────────────┘
54/// ```
55///
56/// The `values` of this `FixedSizeListArray`s are stored in a child
57/// [`StringArray`] where logical null values take up `values_length` slots in the array
58/// as shown in the following diagram. The logical values
59/// are shown on the left, and the actual `FixedSizeListArray` encoding on the right
60///
61/// ```text
62///                                 ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐
63///                                                         ┌ ─ ─ ─ ─ ─ ─ ─ ─┐
64///  ┌─────────────┐                │     ┌───┐               ┌───┐ ┌──────┐      │
65///  │   [A,B]     │                      │ 1 │             │ │ 1 │ │  A   │ │ 0
66///  ├─────────────┤                │     ├───┤               ├───┤ ├──────┤      │
67///  │    NULL     │                      │ 0 │             │ │ 1 │ │  B   │ │ 1
68///  ├─────────────┤                │     ├───┤               ├───┤ ├──────┤      │
69///  │  [C,NULL]   │                      │ 1 │             │ │ 0 │ │ ???? │ │ 2
70///  └─────────────┘                │     └───┘               ├───┤ ├──────┤      │
71///                                                         | │ 0 │ │ ???? │ │ 3
72///  Logical Values                 │   Validity              ├───┤ ├──────┤      │
73///                                     (nulls)             │ │ 1 │ │  C   │ │ 4
74///                                 │                         ├───┤ ├──────┤      │
75///                                                         │ │ 0 │ │ ???? │ │ 5
76///                                 │                         └───┘ └──────┘      │
77///                                                         │     Values     │
78///                                 │   FixedSizeListArray        (Array)         │
79///                                                         └ ─ ─ ─ ─ ─ ─ ─ ─┘
80///                                 └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
81/// ```
82///
83/// # Example
84///
85/// ```
86/// # use std::sync::Arc;
87/// # use arrow_array::{Array, FixedSizeListArray, Int32Array};
88/// # use arrow_data::ArrayData;
89/// # use arrow_schema::{DataType, Field};
90/// # use arrow_buffer::Buffer;
91/// // Construct a value array
92/// let value_data = ArrayData::builder(DataType::Int32)
93///     .len(9)
94///     .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8]))
95///     .build()
96///     .unwrap();
97/// let list_data_type = DataType::FixedSizeList(
98///     Arc::new(Field::new_list_field(DataType::Int32, false)),
99///     3,
100/// );
101/// let list_data = ArrayData::builder(list_data_type.clone())
102///     .len(3)
103///     .add_child_data(value_data.clone())
104///     .build()
105///     .unwrap();
106/// let list_array = FixedSizeListArray::from(list_data);
107/// let list0 = list_array.value(0);
108/// let list1 = list_array.value(1);
109/// let list2 = list_array.value(2);
110///
111/// assert_eq!( &[0, 1, 2], list0.as_any().downcast_ref::<Int32Array>().unwrap().values());
112/// assert_eq!( &[3, 4, 5], list1.as_any().downcast_ref::<Int32Array>().unwrap().values());
113/// assert_eq!( &[6, 7, 8], list2.as_any().downcast_ref::<Int32Array>().unwrap().values());
114/// ```
115///
116/// [`StringArray`]: crate::array::StringArray
117/// [fixed length lists]: https://arrow.apache.org/docs/format/Columnar.html#fixed-size-list-layout
118#[derive(Clone)]
119pub struct FixedSizeListArray {
120    data_type: DataType, // Must be DataType::FixedSizeList(value_length)
121    values: ArrayRef,
122    nulls: Option<NullBuffer>,
123    value_length: i32,
124    len: usize,
125}
126
127impl FixedSizeListArray {
128    /// Create a new [`FixedSizeListArray`] with `size` element size, panicking on failure.
129    ///
130    /// Note that if `size == 0` and `nulls` is `None` (a degenerate, non-nullable
131    /// `FixedSizeListArray`), this function will set the length of the array to 0.
132    ///
133    /// If you would like to have a degenerate, non-nullable `FixedSizeListArray` with arbitrary
134    /// length, use the [`try_new_with_length()`] constructor.
135    ///
136    /// [`try_new_with_length()`]: Self::try_new_with_length
137    ///
138    /// # Panics
139    ///
140    /// Panics if [`Self::try_new`] returns an error
141    pub fn new(field: FieldRef, size: i32, values: ArrayRef, nulls: Option<NullBuffer>) -> Self {
142        Self::try_new(field, size, values, nulls).unwrap()
143    }
144
145    /// Create a new [`FixedSizeListArray`] from the provided parts without validation.
146    ///
147    /// # Safety
148    /// - `size >= 0`
149    /// - `values.len() == len * size as usize`
150    /// - `nulls.len() == len` if `nulls` is `Some`
151    /// - `field.data_type() == values.data_type()`
152    pub unsafe fn new_unchecked(
153        field: FieldRef,
154        size: i32,
155        values: ArrayRef,
156        nulls: Option<NullBuffer>,
157        len: usize,
158    ) -> Self {
159        if cfg!(feature = "force_validate") {
160            return Self::try_new_with_length(field, size, values, nulls, len).unwrap();
161        }
162        Self {
163            data_type: DataType::FixedSizeList(field, size),
164            values,
165            value_length: size,
166            nulls,
167            len,
168        }
169    }
170
171    /// Create a new [`FixedSizeListArray`] from the provided parts, returning an error on failure.
172    ///
173    /// Note that if `size == 0` and `nulls` is `None` (a degenerate, non-nullable
174    /// `FixedSizeListArray`), this function will set the length of the array to 0.
175    ///
176    /// If you would like to have a degenerate, non-nullable `FixedSizeListArray` with arbitrary
177    /// length, use the [`try_new_with_length()`] constructor.
178    ///
179    /// [`try_new_with_length()`]: Self::try_new_with_length
180    ///
181    /// # Errors
182    ///
183    /// * `size < 0`
184    /// * `values.len() != nulls.len() * size` if `nulls` is `Some`
185    /// * `values.data_type() != field.data_type()`
186    /// * `!field.is_nullable() && !nulls.expand(size).contains(values.logical_nulls())`
187    pub fn try_new(
188        field: FieldRef,
189        size: i32,
190        values: ArrayRef,
191        nulls: Option<NullBuffer>,
192    ) -> Result<Self, ArrowError> {
193        let s = size.to_usize().ok_or_else(|| {
194            ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}"))
195        })?;
196
197        if s == 0 {
198            // Note that for degenerate (`size == 0`) and non-nullable `FixedSizeList`s, we will set
199            // the length to 0 (`_or_default`).
200            let len = nulls.as_ref().map(|x| x.len()).unwrap_or_default();
201
202            Self::try_new_with_length(field, size, values, nulls, len)
203        } else {
204            if !values.len().is_multiple_of(s) {
205                return Err(ArrowError::InvalidArgumentError(format!(
206                    "Incorrect length of values buffer for FixedSizeListArray, \
207                     expected a multiple of {s} got {}",
208                    values.len(),
209                )));
210            }
211
212            let len = values.len() / s;
213
214            // Check that the null buffer length is correct (if it exists).
215            if let Some(null_buffer) = &nulls
216                && s * null_buffer.len() != values.len()
217            {
218                return Err(ArrowError::InvalidArgumentError(format!(
219                    "Incorrect length of values buffer for FixedSizeListArray, \
220                            expected {} got {}",
221                    s * null_buffer.len(),
222                    values.len(),
223                )));
224            }
225
226            Self::try_new_with_length(field, size, values, nulls, len)
227        }
228    }
229
230    /// Create a new [`FixedSizeListArray`] from the provided parts, returning an error on failure.
231    ///
232    /// This method exists to allow the construction of arbitrary length degenerate (`size == 0`)
233    /// and non-nullable `FixedSizeListArray`s. If you want a nullable `FixedSizeListArray`, then
234    /// you can use [`try_new()`] instead.
235    ///
236    /// [`try_new()`]: Self::try_new
237    ///
238    /// # Errors
239    ///
240    /// * `size < 0`
241    /// * `nulls.len() != len` if `nulls` is `Some`
242    /// * `values.len() != len * size`
243    /// * `values.data_type() != field.data_type()`
244    /// * `!field.is_nullable() && !nulls.expand(size).contains(values.logical_nulls())`
245    pub fn try_new_with_length(
246        field: FieldRef,
247        size: i32,
248        values: ArrayRef,
249        nulls: Option<NullBuffer>,
250        len: usize,
251    ) -> Result<Self, ArrowError> {
252        let s = size.to_usize().ok_or_else(|| {
253            ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}"))
254        })?;
255
256        if let Some(null_buffer) = &nulls
257            && null_buffer.len() != len
258        {
259            return Err(ArrowError::InvalidArgumentError(format!(
260                "Invalid null buffer for FixedSizeListArray, expected {len} found {}",
261                null_buffer.len()
262            )));
263        }
264
265        if s == 0 && !values.is_empty() {
266            return Err(ArrowError::InvalidArgumentError(format!(
267                "An degenerate FixedSizeListArray should have no underlying values, found {} values",
268                values.len()
269            )));
270        }
271
272        if values.len() != len * s {
273            return Err(ArrowError::InvalidArgumentError(format!(
274                "Incorrect length of values buffer for FixedSizeListArray, expected {} got {}",
275                len * s,
276                values.len(),
277            )));
278        }
279
280        if field.data_type() != values.data_type() {
281            return Err(ArrowError::InvalidArgumentError(format!(
282                "FixedSizeListArray expected data type {} got {} for {:?}",
283                field.data_type(),
284                values.data_type(),
285                field.name()
286            )));
287        }
288
289        if let Some(a) = values.logical_nulls() {
290            let nulls_valid = field.is_nullable()
291                || nulls
292                    .as_ref()
293                    .is_some_and(|n| n.expand(size as _).contains(&a))
294                || (nulls.is_none() && a.null_count() == 0);
295
296            if !nulls_valid {
297                return Err(ArrowError::InvalidArgumentError(format!(
298                    "Found unmasked nulls for non-nullable FixedSizeListArray field {:?}",
299                    field.name()
300                )));
301            }
302        }
303
304        let data_type = DataType::FixedSizeList(field, size);
305        Ok(Self {
306            data_type,
307            values,
308            value_length: size,
309            nulls,
310            len,
311        })
312    }
313
314    /// Create a new [`FixedSizeListArray`] of length `len` where all values are null
315    ///
316    /// # Panics
317    ///
318    /// Panics if
319    ///
320    /// * `size < 0`
321    /// * `size * len` would overflow `usize`
322    pub fn new_null(field: FieldRef, size: i32, len: usize) -> Self {
323        let capacity = size.to_usize().unwrap().checked_mul(len).unwrap();
324        Self {
325            values: make_array(ArrayData::new_null(field.data_type(), capacity)),
326            data_type: DataType::FixedSizeList(field, size),
327            nulls: Some(NullBuffer::new_null(len)),
328            value_length: size,
329            len,
330        }
331    }
332
333    /// Deconstruct this array into its constituent parts
334    pub fn into_parts(self) -> (FieldRef, i32, ArrayRef, Option<NullBuffer>) {
335        let DataType::FixedSizeList(f, _) = self.data_type else {
336            unreachable!()
337        };
338        (f, self.value_length, self.values, self.nulls)
339    }
340
341    /// The field that describes the values of this list.
342    pub fn value_field(&self) -> &FieldRef {
343        match &self.data_type {
344            DataType::FixedSizeList(f, _) => f,
345            _ => unreachable!(),
346        }
347    }
348
349    /// Returns a reference to the values of this list.
350    pub fn values(&self) -> &ArrayRef {
351        &self.values
352    }
353
354    /// Returns a clone of the value type of this list.
355    pub fn value_type(&self) -> DataType {
356        self.values.data_type().clone()
357    }
358
359    /// Returns ith value of this list array.
360    ///
361    /// Note: This method does not check for nulls and the value is arbitrary
362    /// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
363    ///
364    /// # Panics
365    /// Panics if index `i` is out of bounds
366    pub fn value(&self, i: usize) -> ArrayRef {
367        self.values
368            .slice(self.value_offset_at(i), self.value_length() as usize)
369    }
370
371    /// Returns the offset for value at index `i`.
372    ///
373    /// Note this doesn't do any bound checking, for performance reason.
374    ///
375    /// Past `i32::MAX` child elements the offset does not fit the return type
376    /// and wraps negative. Use [`Self::value_offset_at`], which stays in
377    /// `usize`.
378    #[deprecated(since = "60.0.0", note = "Use value_offset_at instead")]
379    #[inline]
380    pub fn value_offset(&self, i: usize) -> i32 {
381        self.value_offset_at(i) as i32
382    }
383
384    /// Returns the length for an element.
385    ///
386    /// All elements have the same length as the array is a fixed size.
387    #[inline]
388    pub const fn value_length(&self) -> i32 {
389        self.value_length
390    }
391
392    /// Returns the offset in the child array of the value at index `i`, in
393    /// elements.
394    ///
395    /// Note this doesn't do any bound checking, for performance reason.
396    #[inline]
397    pub const fn value_offset_at(&self, i: usize) -> usize {
398        i * self.value_length as usize
399    }
400
401    /// Returns a zero-copy slice of this array with the indicated offset and length.
402    ///
403    /// # Panics
404    /// Panics if `offset + len > self.len()`
405    pub fn slice(&self, offset: usize, len: usize) -> Self {
406        assert!(
407            offset.saturating_add(len) <= self.len,
408            "the length + offset of the sliced FixedSizeListArray cannot exceed the existing length"
409        );
410        let size = self.value_length as usize;
411
412        Self {
413            data_type: self.data_type.clone(),
414            values: self.values.slice(offset * size, len * size),
415            nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)),
416            value_length: self.value_length,
417            len,
418        }
419    }
420
421    /// Creates a [`FixedSizeListArray`] from an iterator of primitive values
422    /// # Example
423    /// ```
424    /// # use arrow_array::FixedSizeListArray;
425    /// # use arrow_array::types::Int32Type;
426    ///
427    /// let data = vec![
428    ///    Some(vec![Some(0), Some(1), Some(2)]),
429    ///    None,
430    ///    Some(vec![Some(3), None, Some(5)]),
431    ///    Some(vec![Some(6), Some(7), Some(45)]),
432    /// ];
433    /// let list_array = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(data, 3);
434    /// println!("{:?}", list_array);
435    /// ```
436    pub fn from_iter_primitive<T, P, I>(iter: I, length: i32) -> Self
437    where
438        T: ArrowPrimitiveType,
439        P: IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
440        I: IntoIterator<Item = Option<P>>,
441    {
442        let l = length as usize;
443        let iter = iter.into_iter();
444        let size_hint = iter.size_hint().0;
445        let mut builder = FixedSizeListBuilder::with_capacity(
446            PrimitiveBuilder::<T>::with_capacity(size_hint * l),
447            length,
448            size_hint,
449        );
450
451        for i in iter {
452            match i {
453                Some(p) => {
454                    for t in p {
455                        builder.values().append_option(t);
456                    }
457                    builder.append(true);
458                }
459                None => {
460                    builder.values().append_nulls(l);
461                    builder.append(false)
462                }
463            }
464        }
465        builder.finish()
466    }
467
468    /// constructs a new iterator
469    pub fn iter(&self) -> FixedSizeListIter<'_> {
470        FixedSizeListIter::new(self)
471    }
472}
473
474impl<'a> IntoIterator for &'a FixedSizeListArray {
475    type Item = Option<ArrayRef>;
476    type IntoIter = FixedSizeListIter<'a>;
477
478    fn into_iter(self) -> Self::IntoIter {
479        FixedSizeListIter::new(self)
480    }
481}
482
483impl From<ArrayData> for FixedSizeListArray {
484    fn from(data: ArrayData) -> Self {
485        let (data_type, len, nulls, offset, _buffers, child_data) = data.into_parts();
486
487        let value_length = match data_type {
488            DataType::FixedSizeList(_, len) => len,
489            data_type => {
490                panic!(
491                    "FixedSizeListArray data should contain a FixedSizeList data type, got {data_type}"
492                )
493            }
494        };
495
496        let size = value_length as usize;
497        let values = make_array(child_data[0].slice(offset * size, len * size));
498        Self {
499            data_type,
500            values,
501            nulls,
502            value_length,
503            len,
504        }
505    }
506}
507
508impl From<FixedSizeListArray> for ArrayData {
509    fn from(array: FixedSizeListArray) -> Self {
510        let builder = ArrayDataBuilder::new(array.data_type)
511            .len(array.len)
512            .nulls(array.nulls)
513            .child_data(vec![array.values.to_data()]);
514
515        unsafe { builder.build_unchecked() }
516    }
517}
518
519/// SAFETY: Correctly implements the contract of Arrow Arrays
520unsafe impl Array for FixedSizeListArray {
521    fn as_any(&self) -> &dyn Any {
522        self
523    }
524
525    fn to_data(&self) -> ArrayData {
526        self.clone().into()
527    }
528
529    fn into_data(self) -> ArrayData {
530        self.into()
531    }
532
533    fn data_type(&self) -> &DataType {
534        &self.data_type
535    }
536
537    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
538        Arc::new(self.slice(offset, length))
539    }
540
541    fn len(&self) -> usize {
542        self.len
543    }
544
545    fn is_empty(&self) -> bool {
546        self.len == 0
547    }
548
549    fn shrink_to_fit(&mut self) {
550        self.values.shrink_to_fit();
551        if let Some(nulls) = &mut self.nulls {
552            nulls.shrink_to_fit();
553        }
554    }
555
556    fn offset(&self) -> usize {
557        0
558    }
559
560    fn nulls(&self) -> Option<&NullBuffer> {
561        self.nulls.as_ref()
562    }
563
564    fn logical_null_count(&self) -> usize {
565        // More efficient that the default implementation
566        self.null_count()
567    }
568
569    fn get_buffer_memory_size(&self) -> usize {
570        let mut size = self.values.get_buffer_memory_size();
571        if let Some(n) = self.nulls.as_ref() {
572            size += n.buffer().capacity();
573        }
574        size
575    }
576
577    fn get_array_memory_size(&self) -> usize {
578        let mut size = std::mem::size_of::<Self>() + self.values.get_array_memory_size();
579        if let Some(n) = self.nulls.as_ref() {
580            size += n.buffer().capacity();
581        }
582        size
583    }
584
585    #[cfg(feature = "pool")]
586    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
587        self.values.claim(pool);
588        if let Some(nulls) = &self.nulls {
589            nulls.claim(pool);
590        }
591    }
592}
593
594impl super::ListLikeArray for FixedSizeListArray {
595    fn values(&self) -> &ArrayRef {
596        self.values()
597    }
598
599    fn element_range(&self, index: usize) -> std::ops::Range<usize> {
600        let value_length = self.value_length().as_usize();
601        let offset = index * value_length;
602        offset..(offset + value_length)
603    }
604}
605
606impl ArrayAccessor for FixedSizeListArray {
607    type Item = ArrayRef;
608
609    fn value(&self, index: usize) -> Self::Item {
610        FixedSizeListArray::value(self, index)
611    }
612
613    unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
614        FixedSizeListArray::value(self, index)
615    }
616}
617
618impl std::fmt::Debug for FixedSizeListArray {
619    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
620        write!(f, "FixedSizeListArray<{}>\n[\n", self.value_length())?;
621        print_long_array(self, f, &mut |index, f| {
622            std::fmt::Debug::fmt(&self.value(index), f)
623        })?;
624        write!(f, "]")
625    }
626}
627
628impl ArrayAccessor for &FixedSizeListArray {
629    type Item = ArrayRef;
630
631    fn value(&self, index: usize) -> Self::Item {
632        FixedSizeListArray::value(self, index)
633    }
634
635    unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
636        FixedSizeListArray::value(self, index)
637    }
638}
639
640#[cfg(test)]
641mod tests {
642    use arrow_buffer::{BooleanBuffer, Buffer, bit_util};
643    use arrow_schema::Field;
644
645    use crate::cast::AsArray;
646    use crate::types::Int32Type;
647    use crate::{Int32Array, new_empty_array};
648
649    use super::*;
650
651    #[test]
652    fn test_fixed_size_list_array() {
653        // Construct a value array
654        let value_data = ArrayData::builder(DataType::Int32)
655            .len(9)
656            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8]))
657            .build()
658            .unwrap();
659
660        // Construct a list array from the above two
661        let list_data_type =
662            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
663        let list_data = ArrayData::builder(list_data_type.clone())
664            .len(3)
665            .add_child_data(value_data.clone())
666            .build()
667            .unwrap();
668        let list_array = FixedSizeListArray::from(list_data);
669
670        assert_eq!(value_data, list_array.values().to_data());
671        assert_eq!(DataType::Int32, list_array.value_type());
672        assert_eq!(3, list_array.len());
673        assert_eq!(0, list_array.null_count());
674        assert_eq!(6, list_array.value_offset_at(2));
675        assert_eq!(3, list_array.value_length());
676        assert_eq!(0, list_array.value(0).as_primitive::<Int32Type>().value(0));
677        for i in 0..3 {
678            assert!(list_array.is_valid(i));
679            assert!(!list_array.is_null(i));
680        }
681
682        // Now test with a non-zero offset
683        let list_data = ArrayData::builder(list_data_type)
684            .len(2)
685            .offset(1)
686            .add_child_data(value_data.clone())
687            .build()
688            .unwrap();
689        let list_array = FixedSizeListArray::from(list_data);
690
691        assert_eq!(value_data.slice(3, 6), list_array.values().to_data());
692        assert_eq!(DataType::Int32, list_array.value_type());
693        assert_eq!(2, list_array.len());
694        assert_eq!(0, list_array.null_count());
695        assert_eq!(3, list_array.value(0).as_primitive::<Int32Type>().value(0));
696        assert_eq!(3, list_array.value_offset_at(1));
697        assert_eq!(3, list_array.value_length());
698    }
699
700    #[test]
701    #[should_panic(expected = "assertion failed: end <= self.len()")]
702    // Different error messages, so skip for now
703    // https://github.com/apache/arrow-rs/issues/1545
704    #[cfg(not(feature = "force_validate"))]
705    fn test_fixed_size_list_array_unequal_children() {
706        // Construct a value array
707        let value_data = ArrayData::builder(DataType::Int32)
708            .len(8)
709            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
710            .build()
711            .unwrap();
712
713        // Construct a list array from the above two
714        let list_data_type =
715            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
716        let list_data = unsafe {
717            ArrayData::builder(list_data_type)
718                .len(3)
719                .add_child_data(value_data)
720                .build_unchecked()
721        };
722        drop(FixedSizeListArray::from(list_data));
723    }
724
725    #[test]
726    fn test_fixed_size_list_array_slice() {
727        // Construct a value array
728        let value_data = ArrayData::builder(DataType::Int32)
729            .len(10)
730            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
731            .build()
732            .unwrap();
733
734        // Set null buts for the nested array:
735        //  [[0, 1], null, null, [6, 7], [8, 9]]
736        // 01011001 00000001
737        let mut null_bits: [u8; 1] = [0; 1];
738        bit_util::set_bit(&mut null_bits, 0);
739        bit_util::set_bit(&mut null_bits, 3);
740        bit_util::set_bit(&mut null_bits, 4);
741
742        // Construct a fixed size list array from the above two
743        let list_data_type =
744            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 2);
745        let list_data = ArrayData::builder(list_data_type)
746            .len(5)
747            .add_child_data(value_data.clone())
748            .null_bit_buffer(Some(Buffer::from(null_bits)))
749            .build()
750            .unwrap();
751        let list_array = FixedSizeListArray::from(list_data);
752
753        assert_eq!(value_data, list_array.values().to_data());
754        assert_eq!(DataType::Int32, list_array.value_type());
755        assert_eq!(5, list_array.len());
756        assert_eq!(2, list_array.null_count());
757        assert_eq!(6, list_array.value_offset_at(3));
758        assert_eq!(2, list_array.value_length());
759
760        let sliced_array = list_array.slice(1, 4);
761        assert_eq!(4, sliced_array.len());
762        assert_eq!(2, sliced_array.null_count());
763
764        for i in 0..sliced_array.len() {
765            if bit_util::get_bit(&null_bits, 1 + i) {
766                assert!(sliced_array.is_valid(i));
767            } else {
768                assert!(sliced_array.is_null(i));
769            }
770        }
771
772        // Check offset and length for each non-null value.
773        let sliced_list_array = sliced_array
774            .as_any()
775            .downcast_ref::<FixedSizeListArray>()
776            .unwrap();
777        assert_eq!(2, sliced_list_array.value_length());
778        assert_eq!(4, sliced_list_array.value_offset_at(2));
779        assert_eq!(6, sliced_list_array.value_offset_at(3));
780    }
781
782    #[test]
783    #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
784    fn test_fixed_size_list_array_index_out_of_bound() {
785        // Construct a value array
786        let value_data = ArrayData::builder(DataType::Int32)
787            .len(10)
788            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
789            .build()
790            .unwrap();
791
792        // Set null buts for the nested array:
793        //  [[0, 1], null, null, [6, 7], [8, 9]]
794        // 01011001 00000001
795        let mut null_bits: [u8; 1] = [0; 1];
796        bit_util::set_bit(&mut null_bits, 0);
797        bit_util::set_bit(&mut null_bits, 3);
798        bit_util::set_bit(&mut null_bits, 4);
799
800        // Construct a fixed size list array from the above two
801        let list_data_type =
802            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 2);
803        let list_data = ArrayData::builder(list_data_type)
804            .len(5)
805            .add_child_data(value_data)
806            .null_bit_buffer(Some(Buffer::from(null_bits)))
807            .build()
808            .unwrap();
809        let list_array = FixedSizeListArray::from(list_data);
810
811        list_array.value(10);
812    }
813
814    #[test]
815    fn test_fixed_size_list_constructors() {
816        let values = Arc::new(Int32Array::from_iter([
817            Some(1),
818            Some(2),
819            None,
820            None,
821            Some(3),
822            Some(4),
823        ]));
824
825        let field = Arc::new(Field::new_list_field(DataType::Int32, true));
826        let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), None);
827        assert_eq!(list.len(), 3);
828
829        let nulls = NullBuffer::new_null(3);
830        let list = FixedSizeListArray::new(field.clone(), 2, values.clone(), Some(nulls));
831        assert_eq!(list.len(), 3);
832
833        let list = FixedSizeListArray::new(field.clone(), 3, values.clone(), None);
834        assert_eq!(list.len(), 2);
835
836        let err = FixedSizeListArray::try_new(field.clone(), 4, values.clone(), None).unwrap_err();
837        assert_eq!(
838            err.to_string(),
839            "Invalid argument error: Incorrect length of values buffer for FixedSizeListArray, \
840             expected a multiple of 4 got 6",
841        );
842
843        let err =
844            FixedSizeListArray::try_new_with_length(field.clone(), 4, values.clone(), None, 1)
845                .unwrap_err();
846        assert_eq!(
847            err.to_string(),
848            "Invalid argument error: Incorrect length of values buffer for FixedSizeListArray, expected 4 got 6"
849        );
850
851        let err = FixedSizeListArray::try_new(field.clone(), -1, values.clone(), None).unwrap_err();
852        assert_eq!(
853            err.to_string(),
854            "Invalid argument error: Size cannot be negative, got -1"
855        );
856
857        let nulls = NullBuffer::new_null(2);
858        let err = FixedSizeListArray::try_new(field, 2, values.clone(), Some(nulls)).unwrap_err();
859        assert_eq!(
860            err.to_string(),
861            "Invalid argument error: Incorrect length of values buffer for FixedSizeListArray, expected 4 got 6"
862        );
863
864        let field = Arc::new(Field::new_list_field(DataType::Int32, false));
865        let err = FixedSizeListArray::try_new(field.clone(), 2, values.clone(), None).unwrap_err();
866        assert_eq!(
867            err.to_string(),
868            "Invalid argument error: Found unmasked nulls for non-nullable FixedSizeListArray field \"item\""
869        );
870
871        // Valid as nulls in child masked by parent
872        let nulls = NullBuffer::new(BooleanBuffer::new(Buffer::from([0b0000101]), 0, 3));
873        FixedSizeListArray::new(field, 2, values.clone(), Some(nulls));
874
875        let field = Arc::new(Field::new_list_field(DataType::Int64, true));
876        let err = FixedSizeListArray::try_new(field, 2, values, None).unwrap_err();
877        assert_eq!(
878            err.to_string(),
879            "Invalid argument error: FixedSizeListArray expected data type Int64 got Int32 for \"item\""
880        );
881    }
882
883    #[test]
884    fn degenerate_fixed_size_list() {
885        let field = Arc::new(Field::new_list_field(DataType::Int32, true));
886        let nulls = NullBuffer::new_null(2);
887        let values = new_empty_array(&DataType::Int32);
888        let list = FixedSizeListArray::new(field.clone(), 0, values.clone(), Some(nulls.clone()));
889        assert_eq!(list.len(), 2);
890
891        // Test invalid null buffer length.
892        let err = FixedSizeListArray::try_new_with_length(
893            field.clone(),
894            0,
895            values.clone(),
896            Some(nulls),
897            5,
898        )
899        .unwrap_err();
900        assert_eq!(
901            err.to_string(),
902            "Invalid argument error: Invalid null buffer for FixedSizeListArray, expected 5 found 2"
903        );
904
905        // Test non-empty values for degenerate list.
906        let non_empty_values = Arc::new(Int32Array::from(vec![1, 2, 3]));
907        let err =
908            FixedSizeListArray::try_new_with_length(field.clone(), 0, non_empty_values, None, 3)
909                .unwrap_err();
910        assert_eq!(
911            err.to_string(),
912            "Invalid argument error: An degenerate FixedSizeListArray should have no underlying values, found 3 values"
913        );
914    }
915
916    #[test]
917    fn test_fixed_size_list_new_null_len() {
918        let field = Arc::new(Field::new_list_field(DataType::Int32, true));
919        let array = FixedSizeListArray::new_null(field, 2, 5);
920        assert_eq!(array.len(), 5);
921    }
922}
arrow_array/array/fixed_size_list_array.rs

arrow_array/array/
fixed_size_list_array.rs
