arrow_array/array/
struct_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::{Array, ArrayRef, RecordBatch, make_array, new_null_array};
20use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer};
21use arrow_data::{ArrayData, ArrayDataBuilder};
22use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields};
23use std::sync::Arc;
24use std::{any::Any, ops::Index};
25
26/// An array of [structs](https://arrow.apache.org/docs/format/Columnar.html#struct-layout)
27///
28/// Each child (called *field*) is represented by a separate array.
29///
30/// # Comparison with [RecordBatch]
31///
32/// Both [`RecordBatch`] and [`StructArray`] represent a collection of columns / arrays with the
33/// same length.
34///
35/// However, there are a couple of key differences:
36///
37/// * [`StructArray`] can be nested within other [`Array`], including itself
38/// * [`RecordBatch`] can contain top-level metadata on its associated [`Schema`][arrow_schema::Schema]
39/// * [`StructArray`] can contain top-level nulls, i.e. `null`
40/// * [`RecordBatch`] can only represent nulls in its child columns, i.e. `{"field": null}`
41///
42/// [`StructArray`] is therefore a more general data container than [`RecordBatch`], and as such
43/// code that needs to handle both will typically share an implementation in terms of
44/// [`StructArray`] and convert to/from [`RecordBatch`] as necessary.
45///
46/// [`From`] implementations are provided to facilitate this conversion, however, converting
47/// from a [`StructArray`] containing top-level nulls to a [`RecordBatch`] will panic, as there
48/// is no way to preserve them.
49///
50/// # Example: Create an array from a vector of fields
51///
52/// ```
53/// use std::sync::Arc;
54/// use arrow_array::{Array, ArrayRef, BooleanArray, Int32Array, StructArray};
55/// use arrow_schema::{DataType, Field};
56///
57/// let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
58/// let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
59///
60/// let struct_array = StructArray::from(vec![
61///     (
62///         Arc::new(Field::new("b", DataType::Boolean, false)),
63///         boolean.clone() as ArrayRef,
64///     ),
65///     (
66///         Arc::new(Field::new("c", DataType::Int32, false)),
67///         int.clone() as ArrayRef,
68///     ),
69/// ]);
70/// assert_eq!(struct_array.column(0).as_ref(), boolean.as_ref());
71/// assert_eq!(struct_array.column(1).as_ref(), int.as_ref());
72/// assert_eq!(4, struct_array.len());
73/// assert_eq!(0, struct_array.null_count());
74/// assert_eq!(0, struct_array.offset());
75/// ```
76#[derive(Clone)]
77pub struct StructArray {
78    len: usize,
79    data_type: DataType,
80    nulls: Option<NullBuffer>,
81    fields: Vec<ArrayRef>,
82}
83
84impl StructArray {
85    /// Create a new [`StructArray`] from the provided parts, panicking on failure
86    ///
87    /// # Panics
88    ///
89    /// Panics if [`Self::try_new`] returns an error
90    pub fn new(fields: Fields, arrays: Vec<ArrayRef>, nulls: Option<NullBuffer>) -> Self {
91        Self::try_new(fields, arrays, nulls).unwrap()
92    }
93
94    /// Create a new [`StructArray`] from the provided parts, returning an error on failure
95    ///
96    /// The length will be inferred from the length of the child arrays.  Returns an error if
97    /// there are no child arrays.  Consider using [`Self::try_new_with_length`] if the length
98    /// is known to avoid this.
99    ///
100    /// # Errors
101    ///
102    /// Errors if
103    ///
104    /// * `fields.len() == 0`
105    /// * Any reason that [`Self::try_new_with_length`] would error
106    pub fn try_new(
107        fields: Fields,
108        arrays: Vec<ArrayRef>,
109        nulls: Option<NullBuffer>,
110    ) -> Result<Self, ArrowError> {
111        let len = arrays.first().map(|x| x.len()).ok_or_else(||ArrowError::InvalidArgumentError("use StructArray::try_new_with_length or StructArray::new_empty_fields to create a struct array with no fields so that the length can be set correctly".to_string()))?;
112
113        Self::try_new_with_length(fields, arrays, nulls, len)
114    }
115
116    /// Create a new [`StructArray`] from the provided parts, returning an error on failure
117    ///
118    /// # Errors
119    ///
120    /// Errors if
121    ///
122    /// * `fields.len() != arrays.len()`
123    /// * `fields[i].data_type() != arrays[i].data_type()`
124    /// * `arrays[i].len() != arrays[j].len()`
125    /// * `arrays[i].len() != nulls.len()`
126    /// * `!fields[i].is_nullable() && !nulls.contains(arrays[i].nulls())`
127    pub fn try_new_with_length(
128        fields: Fields,
129        arrays: Vec<ArrayRef>,
130        nulls: Option<NullBuffer>,
131        len: usize,
132    ) -> Result<Self, ArrowError> {
133        if fields.len() != arrays.len() {
134            return Err(ArrowError::InvalidArgumentError(format!(
135                "Incorrect number of arrays for StructArray fields, expected {} got {}",
136                fields.len(),
137                arrays.len()
138            )));
139        }
140
141        if let Some(n) = nulls.as_ref()
142            && n.len() != len
143        {
144            return Err(ArrowError::InvalidArgumentError(format!(
145                "Incorrect number of nulls for StructArray, expected {len} got {}",
146                n.len(),
147            )));
148        }
149
150        for (f, a) in fields.iter().zip(&arrays) {
151            if f.data_type() != a.data_type() {
152                return Err(ArrowError::InvalidArgumentError(format!(
153                    "Incorrect datatype for StructArray field {:?}, expected {} got {}",
154                    f.name(),
155                    f.data_type(),
156                    a.data_type()
157                )));
158            }
159
160            if a.len() != len {
161                return Err(ArrowError::InvalidArgumentError(format!(
162                    "Incorrect array length for StructArray field {:?}, expected {} got {}",
163                    f.name(),
164                    len,
165                    a.len()
166                )));
167            }
168
169            if !f.is_nullable()
170                && let Some(a) = a.logical_nulls()
171                && nulls.as_ref().is_none_or(|n| !n.contains(&a))
172                && a.null_count() > 0
173            {
174                return Err(ArrowError::InvalidArgumentError(format!(
175                    "Found unmasked nulls for non-nullable StructArray field {:?}",
176                    f.name()
177                )));
178            }
179        }
180
181        Ok(Self {
182            len,
183            data_type: DataType::Struct(fields),
184            nulls: nulls.filter(|n| n.null_count() > 0),
185            fields: arrays,
186        })
187    }
188
189    /// Create a new [`StructArray`] of length `len` where all values are null
190    pub fn new_null(fields: Fields, len: usize) -> Self {
191        let arrays = fields
192            .iter()
193            .map(|f| new_null_array(f.data_type(), len))
194            .collect();
195
196        Self {
197            len,
198            data_type: DataType::Struct(fields),
199            nulls: Some(NullBuffer::new_null(len)),
200            fields: arrays,
201        }
202    }
203
204    /// Create a new [`StructArray`] from the provided parts without validation
205    ///
206    /// The length will be inferred from the length of the child arrays.  Panics if there are no
207    /// child arrays.  Consider using [`Self::new_unchecked_with_length`] if the length is known
208    /// to avoid this.
209    ///
210    /// # Safety
211    ///
212    /// Safe if [`Self::new`] would not panic with the given arguments
213    pub unsafe fn new_unchecked(
214        fields: Fields,
215        arrays: Vec<ArrayRef>,
216        nulls: Option<NullBuffer>,
217    ) -> Self {
218        if cfg!(feature = "force_validate") {
219            return Self::new(fields, arrays, nulls);
220        }
221
222        let len = arrays.first().map(|x| x.len()).expect(
223            "cannot use StructArray::new_unchecked if there are no fields, length is unknown",
224        );
225        Self {
226            len,
227            data_type: DataType::Struct(fields),
228            nulls,
229            fields: arrays,
230        }
231    }
232
233    /// Create a new [`StructArray`] from the provided parts without validation
234    ///
235    /// # Safety
236    ///
237    /// Safe if [`Self::new`] would not panic with the given arguments
238    pub unsafe fn new_unchecked_with_length(
239        fields: Fields,
240        arrays: Vec<ArrayRef>,
241        nulls: Option<NullBuffer>,
242        len: usize,
243    ) -> Self {
244        if cfg!(feature = "force_validate") {
245            return Self::try_new_with_length(fields, arrays, nulls, len).unwrap();
246        }
247
248        Self {
249            len,
250            data_type: DataType::Struct(fields),
251            nulls,
252            fields: arrays,
253        }
254    }
255
256    /// Create a new [`StructArray`] containing no fields
257    ///
258    /// # Panics
259    ///
260    /// If `len != nulls.len()`
261    pub fn new_empty_fields(len: usize, nulls: Option<NullBuffer>) -> Self {
262        if let Some(n) = &nulls {
263            assert_eq!(len, n.len())
264        }
265        Self {
266            len,
267            data_type: DataType::Struct(Fields::empty()),
268            fields: vec![],
269            nulls,
270        }
271    }
272
273    /// Deconstruct this array into its constituent parts
274    pub fn into_parts(self) -> (Fields, Vec<ArrayRef>, Option<NullBuffer>) {
275        let f = match self.data_type {
276            DataType::Struct(f) => f,
277            _ => unreachable!(),
278        };
279        (f, self.fields, self.nulls)
280    }
281
282    /// Returns the field at `pos`.
283    ///
284    /// # Panics
285    /// Panics if `pos` is out of bounds
286    pub fn column(&self, pos: usize) -> &ArrayRef {
287        &self.fields[pos]
288    }
289
290    /// Return the number of fields in this struct array
291    pub fn num_columns(&self) -> usize {
292        self.fields.len()
293    }
294
295    /// Returns the fields of the struct array
296    pub fn columns(&self) -> &[ArrayRef] {
297        &self.fields
298    }
299
300    /// Return field names in this struct array
301    pub fn column_names(&self) -> Vec<&str> {
302        match self.data_type() {
303            DataType::Struct(fields) => fields
304                .iter()
305                .map(|f| f.name().as_str())
306                .collect::<Vec<&str>>(),
307            _ => unreachable!("Struct array's data type is not struct!"),
308        }
309    }
310
311    /// Returns the [`Fields`] of this [`StructArray`]
312    pub fn fields(&self) -> &Fields {
313        match self.data_type() {
314            DataType::Struct(f) => f,
315            _ => unreachable!(),
316        }
317    }
318
319    /// Return child array whose field name equals to column_name
320    ///
321    /// Note: A schema can currently have duplicate field names, in which case
322    /// the first field will always be selected.
323    /// This issue will be addressed in [#9205](https://github.com/apache/arrow-rs/issues/9205)
324    pub fn column_by_name(&self, column_name: &str) -> Option<&ArrayRef> {
325        self.fields()
326            .find(column_name)
327            .map(|(pos, _)| self.column(pos))
328    }
329
330    /// Returns the [`FieldRef`] at `pos`.
331    pub fn field(&self, pos: usize) -> &FieldRef {
332        &self.fields()[pos]
333    }
334
335    /// Return the [`FieldRef`] whose name equals to `field_name`
336    ///
337    /// Note: A schema can currently have duplicate field names, in which case
338    /// the first field will always be selected.
339    /// This issue will be addressed in [#9205](https://github.com/apache/arrow-rs/issues/9205)
340    pub fn field_by_name(&self, field_name: &str) -> Option<&FieldRef> {
341        self.fields().find(field_name).map(|(_, field)| field)
342    }
343
344    /// Returns a zero-copy slice of this array with the indicated offset and length.
345    ///
346    /// # Panics
347    /// Panics if `offset + len > self.len()`
348    pub fn slice(&self, offset: usize, len: usize) -> Self {
349        assert!(
350            offset.saturating_add(len) <= self.len,
351            "the length + offset of the sliced StructArray cannot exceed the existing length"
352        );
353
354        let fields = self.fields.iter().map(|a| a.slice(offset, len)).collect();
355
356        Self {
357            len,
358            data_type: self.data_type.clone(),
359            nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)),
360            fields,
361        }
362    }
363
364    /// Returns the children of this [`StructArray`] with the struct's validity
365    /// bitmap AND'd into each child's validity bitmap.
366    ///
367    /// This ensures that positions where the struct itself is null are also
368    /// null in each returned child array. Fields that were non-nullable are
369    /// marked nullable in the returned [`Fields`] when the struct has nulls.
370    ///
371    /// If the struct has no nulls, children and fields are returned as-is.
372    ///
373    /// This mirrors the semantics of C++ Arrow's `StructArray::Flatten`.
374    ///
375    /// # Example
376    ///
377    /// ```
378    /// # use std::sync::Arc;
379    /// # use arrow_array::{Array, ArrayRef, Int32Array, StructArray};
380    /// # use arrow_buffer::{BooleanBuffer, NullBuffer};
381    /// # use arrow_schema::{DataType, Field, Fields};
382    /// let child = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
383    /// let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
384    /// let sa = StructArray::new(
385    ///     Fields::from(vec![Field::new("a", DataType::Int32, false)]),
386    ///     vec![child],
387    ///     Some(struct_nulls),
388    /// );
389    /// let (fields, columns) = sa.flatten();
390    /// assert!(fields[0].is_nullable());
391    /// assert!(columns[0].is_null(1));
392    /// ```
393    pub fn flatten(&self) -> (Fields, Vec<ArrayRef>) {
394        let schema_fields = self.fields();
395
396        let struct_nulls = match &self.nulls {
397            Some(n) => n,
398            None => return (schema_fields.clone(), self.fields.clone()),
399        };
400
401        let new_fields: Fields = schema_fields
402            .iter()
403            .map(|f| {
404                if f.is_nullable() {
405                    Arc::clone(f)
406                } else {
407                    Arc::new(f.as_ref().clone().with_nullable(true))
408                }
409            })
410            .collect::<Vec<_>>()
411            .into();
412
413        let new_columns = self
414            .fields
415            .iter()
416            .map(|child| {
417                let merged = NullBuffer::union(Some(struct_nulls), child.nulls());
418                // SAFETY: We only make the null buffer more restrictive (adding nulls).
419                // All data buffers and child data remain unchanged.
420                let data = child.to_data().into_builder().nulls(merged);
421                make_array(unsafe { data.build_unchecked() })
422            })
423            .collect();
424
425        (new_fields, new_columns)
426    }
427}
428
429impl From<ArrayData> for StructArray {
430    fn from(data: ArrayData) -> Self {
431        let (data_type, len, nulls, offset, _buffers, child_data) = data.into_parts();
432
433        let parent_offset = offset;
434        let parent_len = len;
435
436        let fields = child_data
437            .into_iter()
438            .map(|cd| {
439                if parent_offset != 0 || parent_len != cd.len() {
440                    make_array(cd.slice(parent_offset, parent_len))
441                } else {
442                    make_array(cd)
443                }
444            })
445            .collect();
446
447        Self {
448            len,
449            data_type,
450            nulls,
451            fields,
452        }
453    }
454}
455
456impl From<StructArray> for ArrayData {
457    fn from(array: StructArray) -> Self {
458        let builder = ArrayDataBuilder::new(array.data_type)
459            .len(array.len)
460            .nulls(array.nulls)
461            .child_data(array.fields.iter().map(|x| x.to_data()).collect());
462
463        unsafe { builder.build_unchecked() }
464    }
465}
466
467impl TryFrom<Vec<(&str, ArrayRef)>> for StructArray {
468    type Error = ArrowError;
469
470    /// builds a StructArray from a vector of names and arrays.
471    fn try_from(values: Vec<(&str, ArrayRef)>) -> Result<Self, ArrowError> {
472        let (fields, arrays): (Vec<_>, _) = values
473            .into_iter()
474            .map(|(name, array)| {
475                (
476                    Field::new(name, array.data_type().clone(), array.is_nullable()),
477                    array,
478                )
479            })
480            .unzip();
481
482        StructArray::try_new(fields.into(), arrays, None)
483    }
484}
485
486/// SAFETY: Correctly implements the contract of Arrow Arrays
487unsafe impl Array for StructArray {
488    fn as_any(&self) -> &dyn Any {
489        self
490    }
491
492    fn to_data(&self) -> ArrayData {
493        self.clone().into()
494    }
495
496    fn into_data(self) -> ArrayData {
497        self.into()
498    }
499
500    fn data_type(&self) -> &DataType {
501        &self.data_type
502    }
503
504    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
505        Arc::new(self.slice(offset, length))
506    }
507
508    fn len(&self) -> usize {
509        self.len
510    }
511
512    fn is_empty(&self) -> bool {
513        self.len == 0
514    }
515
516    fn shrink_to_fit(&mut self) {
517        if let Some(nulls) = &mut self.nulls {
518            nulls.shrink_to_fit();
519        }
520        self.fields.iter_mut().for_each(|n| n.shrink_to_fit());
521    }
522
523    fn offset(&self) -> usize {
524        0
525    }
526
527    fn nulls(&self) -> Option<&NullBuffer> {
528        self.nulls.as_ref()
529    }
530
531    fn logical_null_count(&self) -> usize {
532        // More efficient that the default implementation
533        self.null_count()
534    }
535
536    fn get_buffer_memory_size(&self) -> usize {
537        let mut size = self.fields.iter().map(|a| a.get_buffer_memory_size()).sum();
538        if let Some(n) = self.nulls.as_ref() {
539            size += n.buffer().capacity();
540        }
541        size
542    }
543
544    fn get_array_memory_size(&self) -> usize {
545        let mut size = self.fields.iter().map(|a| a.get_array_memory_size()).sum();
546        size += std::mem::size_of::<Self>();
547        if let Some(n) = self.nulls.as_ref() {
548            size += n.buffer().capacity();
549        }
550        size
551    }
552
553    #[cfg(feature = "pool")]
554    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
555        for field in &self.fields {
556            field.claim(pool);
557        }
558        if let Some(nulls) = &self.nulls {
559            nulls.claim(pool);
560        }
561    }
562}
563
564impl From<Vec<(FieldRef, ArrayRef)>> for StructArray {
565    fn from(v: Vec<(FieldRef, ArrayRef)>) -> Self {
566        let (fields, arrays): (Vec<_>, _) = v.into_iter().unzip();
567        StructArray::new(fields.into(), arrays, None)
568    }
569}
570
571impl std::fmt::Debug for StructArray {
572    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
573        writeln!(f, "StructArray")?;
574        writeln!(f, "-- validity:")?;
575        writeln!(f, "[")?;
576        print_long_array(self, f, |_array, _index, f| write!(f, "valid"))?;
577        writeln!(f, "]\n[")?;
578        for (child_index, name) in self.column_names().iter().enumerate() {
579            let column = self.column(child_index);
580            writeln!(
581                f,
582                "-- child {}: \"{}\" ({:?})",
583                child_index,
584                name,
585                column.data_type()
586            )?;
587            std::fmt::Debug::fmt(column, f)?;
588            writeln!(f)?;
589        }
590        write!(f, "]")
591    }
592}
593
594impl From<(Vec<(FieldRef, ArrayRef)>, Buffer)> for StructArray {
595    fn from(pair: (Vec<(FieldRef, ArrayRef)>, Buffer)) -> Self {
596        let len = pair.0.first().map(|x| x.1.len()).unwrap_or_default();
597        let (fields, arrays): (Vec<_>, Vec<_>) = pair.0.into_iter().unzip();
598        let nulls = NullBuffer::new(BooleanBuffer::new(pair.1, 0, len));
599        Self::new(fields.into(), arrays, Some(nulls))
600    }
601}
602
603impl From<RecordBatch> for StructArray {
604    fn from(value: RecordBatch) -> Self {
605        Self {
606            len: value.num_rows(),
607            data_type: DataType::Struct(value.schema().fields().clone()),
608            nulls: None,
609            fields: value.columns().to_vec(),
610        }
611    }
612}
613
614impl Index<&str> for StructArray {
615    type Output = ArrayRef;
616
617    /// Get a reference to a column's array by name.
618    ///
619    /// Note: A schema can currently have duplicate field names, in which case
620    /// the first field will always be selected.
621    /// This issue will be addressed in [ARROW-11178](https://issues.apache.org/jira/browse/ARROW-11178)
622    ///
623    /// # Panics
624    ///
625    /// Panics if the name is not in the schema.
626    fn index(&self, name: &str) -> &Self::Output {
627        self.column_by_name(name).unwrap()
628    }
629}
630
631#[cfg(test)]
632mod tests {
633    use super::*;
634
635    use crate::{BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, StringArray};
636    use arrow_buffer::ToByteSlice;
637
638    #[test]
639    fn test_struct_array_builder() {
640        let boolean_array = BooleanArray::from(vec![false, false, true, true]);
641        let int_array = Int64Array::from(vec![42, 28, 19, 31]);
642
643        let fields = vec![
644            Field::new("a", DataType::Boolean, false),
645            Field::new("b", DataType::Int64, false),
646        ];
647        let struct_array_data = ArrayData::builder(DataType::Struct(fields.into()))
648            .len(4)
649            .add_child_data(boolean_array.to_data())
650            .add_child_data(int_array.to_data())
651            .build()
652            .unwrap();
653        let struct_array = StructArray::from(struct_array_data);
654
655        assert_eq!(struct_array.column(0).as_ref(), &boolean_array);
656        assert_eq!(struct_array.column(1).as_ref(), &int_array);
657    }
658
659    #[test]
660    fn test_struct_array_from() {
661        let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
662        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
663
664        let struct_array = StructArray::from(vec![
665            (
666                Arc::new(Field::new("b", DataType::Boolean, false)),
667                boolean.clone() as ArrayRef,
668            ),
669            (
670                Arc::new(Field::new("c", DataType::Int32, false)),
671                int.clone() as ArrayRef,
672            ),
673        ]);
674        assert_eq!(struct_array.column(0).as_ref(), boolean.as_ref());
675        assert_eq!(struct_array.column(1).as_ref(), int.as_ref());
676        assert_eq!(4, struct_array.len());
677        assert_eq!(0, struct_array.null_count());
678        assert_eq!(0, struct_array.offset());
679    }
680
681    #[test]
682    fn test_struct_array_from_data_with_offset_and_length() {
683        // Various ways to make the struct array:
684        //
685        // [{x: 2}, {x: 3}, None]
686        //
687        // from slicing larger buffers/arrays with offsets and lengths
688        let int_arr = Int32Array::from(vec![1, 2, 3, 4, 5]);
689        let int_field = Field::new("x", DataType::Int32, false);
690        let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, true, false]));
691        let int_data = int_arr.to_data();
692        // Case 1: Offset + length, nulls are not sliced
693        let case1 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
694            .len(3)
695            .offset(1)
696            .nulls(Some(struct_nulls))
697            .add_child_data(int_data.clone())
698            .build()
699            .unwrap();
700
701        // Case 2: Offset + length, nulls are sliced
702        let struct_nulls =
703            NullBuffer::new(BooleanBuffer::from(vec![true, true, true, false, true]).slice(1, 3));
704        let case2 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
705            .len(3)
706            .offset(1)
707            .nulls(Some(struct_nulls.clone()))
708            .add_child_data(int_data.clone())
709            .build()
710            .unwrap();
711
712        // Case 3: struct length is smaller than child length but no offset
713        let offset_int_data = int_data.slice(1, 4);
714        let case3 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
715            .len(3)
716            .nulls(Some(struct_nulls))
717            .add_child_data(offset_int_data)
718            .build()
719            .unwrap();
720
721        let expected = StructArray::new(
722            Fields::from(vec![int_field.clone()]),
723            vec![Arc::new(int_arr)],
724            Some(NullBuffer::new(BooleanBuffer::from(vec![
725                true, true, true, false, true,
726            ]))),
727        )
728        .slice(1, 3);
729
730        for case in [case1, case2, case3] {
731            let struct_arr_from_data = StructArray::from(case);
732            assert_eq!(struct_arr_from_data, expected);
733            assert_eq!(struct_arr_from_data.column(0), expected.column(0));
734        }
735    }
736
737    #[test]
738    #[should_panic(expected = "assertion failed: end <= self.len()")]
739    fn test_struct_array_from_data_with_offset_and_length_error() {
740        let int_arr = Int32Array::from(vec![1, 2, 3, 4, 5]);
741        let int_field = Field::new("x", DataType::Int32, false);
742        let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, true, false]));
743        let int_data = int_arr.to_data();
744        // If parent offset is 3 and len is 3 then child must have 6 items
745        let struct_data =
746            ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()])))
747                .len(3)
748                .offset(3)
749                .nulls(Some(struct_nulls))
750                .add_child_data(int_data)
751                .build()
752                .unwrap();
753        let _ = StructArray::from(struct_data);
754    }
755
756    /// validates that struct can be accessed using `column_name` as index i.e. `struct_array["column_name"]`.
757    #[test]
758    fn test_struct_array_index_access() {
759        let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
760        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
761
762        let struct_array = StructArray::from(vec![
763            (
764                Arc::new(Field::new("b", DataType::Boolean, false)),
765                boolean.clone() as ArrayRef,
766            ),
767            (
768                Arc::new(Field::new("c", DataType::Int32, false)),
769                int.clone() as ArrayRef,
770            ),
771        ]);
772        assert_eq!(struct_array["b"].as_ref(), boolean.as_ref());
773        assert_eq!(struct_array["c"].as_ref(), int.as_ref());
774    }
775
776    #[test]
777    fn test_struct_array_field_access() {
778        let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
779        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
780
781        let b_field = Arc::new(Field::new("b", DataType::Boolean, false));
782        let c_field = Arc::new(Field::new("c", DataType::Int32, false));
783
784        let struct_array = StructArray::from(vec![
785            (b_field.clone(), boolean as ArrayRef),
786            (c_field.clone(), int as ArrayRef),
787        ]);
788
789        assert_eq!(struct_array.field(0), &b_field);
790        assert_eq!(struct_array.field(1), &c_field);
791
792        assert_eq!(struct_array.field_by_name("b"), Some(&b_field));
793        assert_eq!(struct_array.field_by_name("c"), Some(&c_field));
794        assert_eq!(struct_array.field_by_name("d"), None);
795    }
796
797    /// validates that the in-memory representation follows [the spec](https://arrow.apache.org/docs/format/Columnar.html#struct-layout)
798    #[test]
799    fn test_struct_array_from_vec() {
800        let strings: ArrayRef = Arc::new(StringArray::from(vec![
801            Some("joe"),
802            None,
803            None,
804            Some("mark"),
805        ]));
806        let ints: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)]));
807
808        let arr =
809            StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())]).unwrap();
810
811        let struct_data = arr.into_data();
812        assert_eq!(4, struct_data.len());
813        assert_eq!(0, struct_data.null_count());
814
815        let expected_string_data = ArrayData::builder(DataType::Utf8)
816            .len(4)
817            .null_bit_buffer(Some(Buffer::from(&[9_u8])))
818            .add_buffer(Buffer::from([0, 3, 3, 3, 7].to_byte_slice()))
819            .add_buffer(Buffer::from(b"joemark"))
820            .build()
821            .unwrap();
822
823        let expected_int_data = ArrayData::builder(DataType::Int32)
824            .len(4)
825            .null_bit_buffer(Some(Buffer::from(&[11_u8])))
826            .add_buffer(Buffer::from([1, 2, 0, 4].to_byte_slice()))
827            .build()
828            .unwrap();
829
830        assert_eq!(expected_string_data, struct_data.child_data()[0]);
831        assert_eq!(expected_int_data, struct_data.child_data()[1]);
832    }
833
834    #[test]
835    fn test_struct_array_from_vec_error() {
836        let strings: ArrayRef = Arc::new(StringArray::from(vec![
837            Some("joe"),
838            None,
839            None,
840            // 3 elements, not 4
841        ]));
842        let ints: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)]));
843
844        let err = StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
845            .unwrap_err()
846            .to_string();
847
848        assert_eq!(
849            err,
850            "Invalid argument error: Incorrect array length for StructArray field \"f2\", expected 3 got 4"
851        )
852    }
853
854    #[test]
855    #[should_panic(
856        expected = "Incorrect datatype for StructArray field \\\"b\\\", expected Int16 got Boolean"
857    )]
858    fn test_struct_array_from_mismatched_types_single() {
859        drop(StructArray::from(vec![(
860            Arc::new(Field::new("b", DataType::Int16, false)),
861            Arc::new(BooleanArray::from(vec![false, false, true, true])) as Arc<dyn Array>,
862        )]));
863    }
864
865    #[test]
866    #[should_panic(
867        expected = "Incorrect datatype for StructArray field \\\"b\\\", expected Int16 got Boolean"
868    )]
869    fn test_struct_array_from_mismatched_types_multiple() {
870        drop(StructArray::from(vec![
871            (
872                Arc::new(Field::new("b", DataType::Int16, false)),
873                Arc::new(BooleanArray::from(vec![false, false, true, true])) as Arc<dyn Array>,
874            ),
875            (
876                Arc::new(Field::new("c", DataType::Utf8, false)),
877                Arc::new(Int32Array::from(vec![42, 28, 19, 31])),
878            ),
879        ]));
880    }
881
882    #[test]
883    fn test_struct_array_slice() {
884        let boolean_data = ArrayData::builder(DataType::Boolean)
885            .len(5)
886            .add_buffer(Buffer::from([0b00010000]))
887            .null_bit_buffer(Some(Buffer::from([0b00010001])))
888            .build()
889            .unwrap();
890        let int_data = ArrayData::builder(DataType::Int32)
891            .len(5)
892            .add_buffer(Buffer::from([0, 28, 42, 0, 0].to_byte_slice()))
893            .null_bit_buffer(Some(Buffer::from([0b00000110])))
894            .build()
895            .unwrap();
896
897        let field_types = vec![
898            Field::new("a", DataType::Boolean, true),
899            Field::new("b", DataType::Int32, true),
900        ];
901        let struct_array_data = ArrayData::builder(DataType::Struct(field_types.into()))
902            .len(5)
903            .add_child_data(boolean_data.clone())
904            .add_child_data(int_data.clone())
905            .null_bit_buffer(Some(Buffer::from([0b00010111])))
906            .build()
907            .unwrap();
908        let struct_array = StructArray::from(struct_array_data);
909
910        assert_eq!(5, struct_array.len());
911        assert_eq!(1, struct_array.null_count());
912        assert!(struct_array.is_valid(0));
913        assert!(struct_array.is_valid(1));
914        assert!(struct_array.is_valid(2));
915        assert!(struct_array.is_null(3));
916        assert!(struct_array.is_valid(4));
917        assert_eq!(boolean_data, struct_array.column(0).to_data());
918        assert_eq!(int_data, struct_array.column(1).to_data());
919
920        let c0 = struct_array.column(0);
921        let c0 = c0.as_any().downcast_ref::<BooleanArray>().unwrap();
922        assert_eq!(5, c0.len());
923        assert_eq!(3, c0.null_count());
924        assert!(c0.is_valid(0));
925        assert!(!c0.value(0));
926        assert!(c0.is_null(1));
927        assert!(c0.is_null(2));
928        assert!(c0.is_null(3));
929        assert!(c0.is_valid(4));
930        assert!(c0.value(4));
931
932        let c1 = struct_array.column(1);
933        let c1 = c1.as_any().downcast_ref::<Int32Array>().unwrap();
934        assert_eq!(5, c1.len());
935        assert_eq!(3, c1.null_count());
936        assert!(c1.is_null(0));
937        assert!(c1.is_valid(1));
938        assert_eq!(28, c1.value(1));
939        assert!(c1.is_valid(2));
940        assert_eq!(42, c1.value(2));
941        assert!(c1.is_null(3));
942        assert!(c1.is_null(4));
943
944        let sliced_array = struct_array.slice(2, 3);
945        let sliced_array = sliced_array.as_any().downcast_ref::<StructArray>().unwrap();
946        assert_eq!(3, sliced_array.len());
947        assert_eq!(1, sliced_array.null_count());
948        assert!(sliced_array.is_valid(0));
949        assert!(sliced_array.is_null(1));
950        assert!(sliced_array.is_valid(2));
951
952        let sliced_c0 = sliced_array.column(0);
953        let sliced_c0 = sliced_c0.as_any().downcast_ref::<BooleanArray>().unwrap();
954        assert_eq!(3, sliced_c0.len());
955        assert!(sliced_c0.is_null(0));
956        assert!(sliced_c0.is_null(1));
957        assert!(sliced_c0.is_valid(2));
958        assert!(sliced_c0.value(2));
959
960        let sliced_c1 = sliced_array.column(1);
961        let sliced_c1 = sliced_c1.as_any().downcast_ref::<Int32Array>().unwrap();
962        assert_eq!(3, sliced_c1.len());
963        assert!(sliced_c1.is_valid(0));
964        assert_eq!(42, sliced_c1.value(0));
965        assert!(sliced_c1.is_null(1));
966        assert!(sliced_c1.is_null(2));
967    }
968
969    #[test]
970    #[should_panic(
971        expected = "Incorrect array length for StructArray field \\\"c\\\", expected 1 got 2"
972    )]
973    fn test_invalid_struct_child_array_lengths() {
974        drop(StructArray::from(vec![
975            (
976                Arc::new(Field::new("b", DataType::Float32, false)),
977                Arc::new(Float32Array::from(vec![1.1])) as Arc<dyn Array>,
978            ),
979            (
980                Arc::new(Field::new("c", DataType::Float64, false)),
981                Arc::new(Float64Array::from(vec![2.2, 3.3])),
982            ),
983        ]));
984    }
985
986    #[test]
987    #[should_panic(expected = "use StructArray::try_new_with_length")]
988    fn test_struct_array_from_empty() {
989        // This can't work because we don't know how many rows the array should have.  Previously we inferred 0 but
990        // that often led to bugs.
991        let _ = StructArray::from(vec![]);
992    }
993
994    #[test]
995    fn test_empty_struct_array() {
996        assert!(StructArray::try_new(Fields::empty(), vec![], None).is_err());
997
998        let arr = StructArray::new_empty_fields(10, None);
999        assert_eq!(arr.len(), 10);
1000        assert_eq!(arr.null_count(), 0);
1001        assert_eq!(arr.num_columns(), 0);
1002
1003        let arr2 = StructArray::try_new_with_length(Fields::empty(), vec![], None, 10).unwrap();
1004        assert_eq!(arr2.len(), 10);
1005
1006        let arr = StructArray::new_empty_fields(10, Some(NullBuffer::new_null(10)));
1007        assert_eq!(arr.len(), 10);
1008        assert_eq!(arr.null_count(), 10);
1009        assert_eq!(arr.num_columns(), 0);
1010
1011        let arr2 = StructArray::try_new_with_length(
1012            Fields::empty(),
1013            vec![],
1014            Some(NullBuffer::new_null(10)),
1015            10,
1016        )
1017        .unwrap();
1018        assert_eq!(arr2.len(), 10);
1019    }
1020
1021    #[test]
1022    #[should_panic(expected = "Found unmasked nulls for non-nullable StructArray field \\\"c\\\"")]
1023    fn test_struct_array_from_mismatched_nullability() {
1024        drop(StructArray::from(vec![(
1025            Arc::new(Field::new("c", DataType::Int32, false)),
1026            Arc::new(Int32Array::from(vec![Some(42), None, Some(19)])) as ArrayRef,
1027        )]));
1028    }
1029
1030    #[test]
1031    fn test_struct_array_fmt_debug() {
1032        let arr: StructArray = StructArray::new(
1033            vec![Arc::new(Field::new("c", DataType::Int32, true))].into(),
1034            vec![Arc::new(Int32Array::from((0..30).collect::<Vec<_>>())) as ArrayRef],
1035            Some(NullBuffer::new(BooleanBuffer::from(
1036                (0..30).map(|i| i % 2 == 0).collect::<Vec<_>>(),
1037            ))),
1038        );
1039        assert_eq!(
1040            format!("{arr:?}"),
1041            "StructArray\n-- validity:\n[\n  valid,\n  null,\n  valid,\n  null,\n  valid,\n  null,\n  valid,\n  null,\n  valid,\n  null,\n  ...10 elements...,\n  valid,\n  null,\n  valid,\n  null,\n  valid,\n  null,\n  valid,\n  null,\n  valid,\n  null,\n]\n[\n-- child 0: \"c\" (Int32)\nPrimitiveArray<Int32>\n[\n  0,\n  1,\n  2,\n  3,\n  4,\n  5,\n  6,\n  7,\n  8,\n  9,\n  ...10 elements...,\n  20,\n  21,\n  22,\n  23,\n  24,\n  25,\n  26,\n  27,\n  28,\n  29,\n]\n]"
1042        )
1043    }
1044
1045    #[test]
1046    fn test_struct_array_logical_nulls() {
1047        // Field is non-nullable
1048        let field = Field::new("a", DataType::Int32, false);
1049        let values = vec![1, 2, 3];
1050        // Create a NullBuffer with all bits set to valid (true)
1051        let nulls = NullBuffer::from(vec![true, true, true]);
1052        let array = Int32Array::new(values.into(), Some(nulls));
1053        let child = Arc::new(array) as ArrayRef;
1054        assert!(child.logical_nulls().is_some());
1055        assert_eq!(child.logical_nulls().unwrap().null_count(), 0);
1056
1057        let fields = Fields::from(vec![field]);
1058        let arrays = vec![child];
1059        let nulls = None;
1060
1061        StructArray::try_new(fields, arrays, nulls).expect("should not error");
1062    }
1063
1064    #[test]
1065    fn test_flatten_no_nulls() {
1066        let child = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
1067        let sa = StructArray::from(vec![(
1068            Arc::new(Field::new("a", DataType::Int32, false)),
1069            child,
1070        )]);
1071
1072        let (fields, columns) = sa.flatten();
1073
1074        assert_eq!(columns.len(), 1);
1075        assert!(!fields[0].is_nullable());
1076        assert_eq!(columns[0].null_count(), 0);
1077        assert_eq!(columns[0].len(), 3);
1078    }
1079
1080    #[test]
1081    fn test_flatten_struct_nulls_child_no_nulls() {
1082        let child = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
1083        let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
1084        let sa = StructArray::new(
1085            Fields::from(vec![Field::new("a", DataType::Int32, false)]),
1086            vec![child],
1087            Some(struct_nulls),
1088        );
1089
1090        let (fields, columns) = sa.flatten();
1091
1092        assert!(fields[0].is_nullable());
1093        assert!(columns[0].is_valid(0));
1094        assert!(columns[0].is_null(1));
1095        assert!(columns[0].is_valid(2));
1096        assert_eq!(columns[0].null_count(), 1);
1097    }
1098
1099    #[test]
1100    fn test_flatten_both_have_nulls() {
1101        // struct validity: [valid, null,  valid, valid]
1102        // child validity:  [valid, valid, null,  valid]
1103        // expected:        [valid, null,  null,  valid]
1104        let child = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])) as ArrayRef;
1105        let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true, true]));
1106        let sa = StructArray::new(
1107            Fields::from(vec![Field::new("a", DataType::Int32, true)]),
1108            vec![child],
1109            Some(struct_nulls),
1110        );
1111
1112        let (fields, columns) = sa.flatten();
1113
1114        assert!(fields[0].is_nullable());
1115        assert!(columns[0].is_valid(0));
1116        assert!(columns[0].is_null(1));
1117        assert!(columns[0].is_null(2));
1118        assert!(columns[0].is_valid(3));
1119        assert_eq!(columns[0].null_count(), 2);
1120    }
1121
1122    #[test]
1123    fn test_flatten_sliced_struct() {
1124        let child = Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef;
1125        let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true, false]));
1126        let sa = StructArray::new(
1127            Fields::from(vec![Field::new("a", DataType::Int32, false)]),
1128            vec![child],
1129            Some(struct_nulls),
1130        );
1131        let sliced = sa.slice(1, 2);
1132
1133        let (fields, columns) = sliced.flatten();
1134
1135        assert!(fields[0].is_nullable());
1136        assert_eq!(columns[0].len(), 2);
1137        assert!(columns[0].is_null(0));
1138        assert!(columns[0].is_valid(1));
1139    }
1140
1141    #[test]
1142    fn test_flatten_multiple_children() {
1143        let int_child = Arc::new(Int32Array::from(vec![Some(1), Some(2), None])) as ArrayRef;
1144        let str_child = Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])) as ArrayRef;
1145        let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
1146        let sa = StructArray::new(
1147            Fields::from(vec![
1148                Field::new("ints", DataType::Int32, true),
1149                Field::new("strs", DataType::Utf8, true),
1150            ]),
1151            vec![int_child, str_child],
1152            Some(struct_nulls),
1153        );
1154
1155        let (fields, columns) = sa.flatten();
1156
1157        assert_eq!(fields.len(), 2);
1158        // int: [valid, null(struct), null(child)] => null_count=2
1159        assert_eq!(columns[0].null_count(), 2);
1160        assert!(columns[0].is_valid(0));
1161        assert!(columns[0].is_null(1));
1162        assert!(columns[0].is_null(2));
1163        // str: [valid, null(struct+child), valid] => null_count=1
1164        assert_eq!(columns[1].null_count(), 1);
1165        assert!(columns[1].is_valid(0));
1166        assert!(columns[1].is_null(1));
1167        assert!(columns[1].is_valid(2));
1168    }
1169
1170    #[test]
1171    fn test_flatten_empty_struct() {
1172        let sa = StructArray::new_empty_fields(5, Some(NullBuffer::new_null(5)));
1173
1174        let (fields, columns) = sa.flatten();
1175
1176        assert_eq!(fields.len(), 0);
1177        assert_eq!(columns.len(), 0);
1178    }
1179
1180    #[test]
1181    fn test_flatten_field_nullability_update() {
1182        let non_null_child = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
1183        let nullable_child = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef;
1184        let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, true, false]));
1185        let sa = StructArray::new(
1186            Fields::from(vec![
1187                Field::new("non_null", DataType::Int32, false),
1188                Field::new("nullable", DataType::Int32, true),
1189            ]),
1190            vec![non_null_child, nullable_child],
1191            Some(struct_nulls),
1192        );
1193
1194        let (fields, _columns) = sa.flatten();
1195
1196        assert!(fields[0].is_nullable()); // was false, now true
1197        assert!(fields[1].is_nullable()); // was true, stays true
1198    }
1199}
arrow_array/array/struct_array.rs

arrow_array/array/
struct_array.rs
