arrow_array/array/
fixed_size_list_array.rs

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

arrow_array/array/
fixed_size_list_array.rs
