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

arrow_array/array/
fixed_size_list_array.rs
