arrow_array/builder/
struct_array_assembler.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::{ArrayRef, StructArray};
19use arrow_buffer::NullBuffer;
20use arrow_schema::{ArrowError, Field, FieldRef, Fields};
21use std::sync::Arc;
22
23/// Assembles a [`StructArray`] from completed child arrays.
24///
25/// Unlike [`StructBuilder`](super::StructBuilder), which incrementally builds
26/// child arrays, this assembler combines arrays that have already been built.
27///
28/// # Example
29///
30/// ```
31/// # use std::sync::Arc;
32/// # use arrow_array::builder::StructArrayAssembler;
33/// # use arrow_array::{Array, ArrayRef, Int32Array, StringArray};
34///
35/// let names = Arc::new(StringArray::from(vec!["one", "two"])) as ArrayRef;
36/// let values = Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef;
37/// // Create struct array with `{name: ..., value: ...}` rows
38/// let array = StructArrayAssembler::new()
39///     .with_field("name", names, false)
40///     .with_field("value", values, false)
41///     .build()
42///     .unwrap();
43///
44/// assert_eq!(array.len(), 2);
45/// ```
46#[derive(Debug, Default, Clone)]
47pub struct StructArrayAssembler {
48    fields: Vec<FieldRef>,
49    arrays: Vec<ArrayRef>,
50    nulls: Option<NullBuffer>,
51}
52
53impl StructArrayAssembler {
54    /// Creates a new empty [`StructArrayAssembler`].
55    pub fn new() -> Self {
56        Self::default()
57    }
58
59    /// Adds an array with a field constructed from its data type.
60    pub fn with_field(
61        mut self,
62        field_name: impl Into<String>,
63        array: ArrayRef,
64        nullable: bool,
65    ) -> Self {
66        self.fields.push(Arc::new(Field::new(
67            field_name,
68            array.data_type().clone(),
69            nullable,
70        )));
71        self.arrays.push(array);
72        self
73    }
74
75    /// Adds an array using a caller-supplied field.
76    ///
77    /// This preserves field metadata that would be lost if the field were
78    /// constructed from the array's data type alone.
79    pub fn with_field_ref(mut self, field: FieldRef, array: ArrayRef) -> Self {
80        self.fields.push(field);
81        self.arrays.push(array);
82        self
83    }
84
85    /// Sets the top-level null buffer for the struct array.
86    pub fn with_nulls(mut self, nulls: NullBuffer) -> Self {
87        self.nulls = Some(nulls);
88        self
89    }
90
91    /// Builds the [`StructArray`].
92    ///
93    /// # Errors
94    ///
95    /// Returns an error if the fields, child arrays, or null buffer have
96    /// incompatible data types or lengths.
97    pub fn build(self) -> Result<StructArray, ArrowError> {
98        StructArray::try_new(Fields::from(self.fields), self.arrays, self.nulls)
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use crate::{Array, Int32Array, StringArray};
106    use arrow_schema::DataType;
107    use std::collections::HashMap;
108
109    #[test]
110    fn build_from_completed_arrays() {
111        let names = Arc::new(StringArray::from(vec!["one", "two"])) as ArrayRef;
112        let values = Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef;
113        let metadata = HashMap::from([("key".to_string(), "value".to_string())]);
114        let value_field =
115            Arc::new(Field::new("value", DataType::Int32, false).with_metadata(metadata.clone()));
116        let nulls = NullBuffer::from(vec![true, false]);
117
118        let array = StructArrayAssembler::new()
119            .with_field("name", names, false)
120            .with_field_ref(value_field, values)
121            .with_nulls(nulls.clone())
122            .build()
123            .unwrap();
124
125        assert_eq!(array.len(), 2);
126        assert_eq!(array.column_names(), &["name", "value"]);
127        assert_eq!(array.fields()[1].metadata(), &metadata);
128        assert_eq!(array.nulls(), Some(&nulls));
129    }
130
131    #[test]
132    fn build_errors_on_mismatched_child_lengths() {
133        let names = Arc::new(StringArray::from(vec!["one", "two"])) as ArrayRef;
134        let values = Arc::new(Int32Array::from(vec![1])) as ArrayRef;
135
136        let err = StructArrayAssembler::new()
137            .with_field("name", names, false)
138            .with_field("value", values, false)
139            .build()
140            .unwrap_err();
141
142        assert_eq!(
143            err.to_string(),
144            "Invalid argument error: Incorrect array length for StructArray field \"value\", expected 2 got 1"
145        );
146    }
147
148    #[test]
149    fn build_errors_on_mismatched_field_data_type() {
150        let value_field = Arc::new(Field::new("value", DataType::Int64, false));
151        let values = Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef;
152
153        let err = StructArrayAssembler::new()
154            .with_field_ref(value_field, values)
155            .build()
156            .unwrap_err();
157
158        assert_eq!(
159            err.to_string(),
160            "Invalid argument error: Incorrect datatype for StructArray field \"value\", expected Int64 got Int32"
161        );
162    }
163}
arrow_array/builder/struct_array_assembler.rs

arrow_array/builder/
struct_array_assembler.rs
