arrow_array/builder/
union_builder.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::builder::buffer_builder::{Int8BufferBuilder, Int32BufferBuilder};
19use crate::builder::{ArrayBuilder, BufferBuilder};
20use crate::{ArrayRef, ArrowPrimitiveType, UnionArray, make_array};
21use arrow_buffer::NullBufferBuilder;
22use arrow_buffer::{ArrowNativeType, Buffer, ScalarBuffer};
23use arrow_data::ArrayDataBuilder;
24use arrow_schema::{ArrowError, DataType, Field};
25use std::any::Any;
26use std::collections::BTreeMap;
27use std::sync::Arc;
28
29/// `FieldData` is a helper struct to track the state of the fields in the `UnionBuilder`.
30#[derive(Debug)]
31struct FieldData {
32    /// The type id for this field
33    type_id: i8,
34    /// The Arrow data type represented in the `values_buffer`, which is untyped
35    data_type: DataType,
36    /// A buffer containing the values for this field in raw bytes
37    values_buffer: Box<dyn FieldDataValues>,
38    ///  The number of array slots represented by the buffer
39    slots: usize,
40    /// A builder for the null bitmap
41    null_buffer_builder: NullBufferBuilder,
42}
43
44/// A type-erased [`BufferBuilder`] used by [`FieldData`]
45trait FieldDataValues: std::fmt::Debug + Send + Sync {
46    fn as_mut_any(&mut self) -> &mut dyn Any;
47
48    fn append_null(&mut self);
49
50    fn finish(&mut self) -> Buffer;
51
52    fn finish_cloned(&self) -> Buffer;
53}
54
55impl<T: ArrowNativeType> FieldDataValues for BufferBuilder<T> {
56    fn as_mut_any(&mut self) -> &mut dyn Any {
57        self
58    }
59
60    fn append_null(&mut self) {
61        self.advance(1)
62    }
63
64    fn finish(&mut self) -> Buffer {
65        self.finish()
66    }
67
68    fn finish_cloned(&self) -> Buffer {
69        Buffer::from_slice_ref(self.as_slice())
70    }
71}
72
73impl FieldData {
74    /// Creates a new `FieldData`.
75    fn new<T: ArrowPrimitiveType>(type_id: i8, data_type: DataType, capacity: usize) -> Self {
76        Self {
77            type_id,
78            data_type,
79            slots: 0,
80            values_buffer: Box::new(BufferBuilder::<T::Native>::new(capacity)),
81            null_buffer_builder: NullBufferBuilder::new(capacity),
82        }
83    }
84
85    /// Appends a single value to this `FieldData`'s `values_buffer`.
86    fn append_value<T: ArrowPrimitiveType>(&mut self, v: T::Native) {
87        self.values_buffer
88            .as_mut_any()
89            .downcast_mut::<BufferBuilder<T::Native>>()
90            .expect("Tried to append unexpected type")
91            .append(v);
92
93        self.null_buffer_builder.append(true);
94        self.slots += 1;
95    }
96
97    /// Appends a null to this `FieldData`.
98    fn append_null(&mut self) {
99        self.values_buffer.append_null();
100        self.null_buffer_builder.append(false);
101        self.slots += 1;
102    }
103}
104
105/// Builder for [`UnionArray`]
106///
107/// Example: **Dense Memory Layout**
108///
109/// ```
110/// # use arrow_array::builder::UnionBuilder;
111/// # use arrow_array::types::{Float64Type, Int32Type};
112///
113/// let mut builder = UnionBuilder::new_dense();
114/// builder.append::<Int32Type>("a", 1).unwrap();
115/// builder.append::<Float64Type>("b", 3.0).unwrap();
116/// builder.append::<Int32Type>("a", 4).unwrap();
117/// let union = builder.build().unwrap();
118///
119/// assert_eq!(union.type_id(0), 0);
120/// assert_eq!(union.type_id(1), 1);
121/// assert_eq!(union.type_id(2), 0);
122///
123/// assert_eq!(union.value_offset(0), 0);
124/// assert_eq!(union.value_offset(1), 0);
125/// assert_eq!(union.value_offset(2), 1);
126/// ```
127///
128/// Example: **Sparse Memory Layout**
129/// ```
130/// # use arrow_array::builder::UnionBuilder;
131/// # use arrow_array::types::{Float64Type, Int32Type};
132///
133/// let mut builder = UnionBuilder::new_sparse();
134/// builder.append::<Int32Type>("a", 1).unwrap();
135/// builder.append::<Float64Type>("b", 3.0).unwrap();
136/// builder.append::<Int32Type>("a", 4).unwrap();
137/// let union = builder.build().unwrap();
138///
139/// assert_eq!(union.type_id(0), 0);
140/// assert_eq!(union.type_id(1), 1);
141/// assert_eq!(union.type_id(2), 0);
142///
143/// assert_eq!(union.value_offset(0), 0);
144/// assert_eq!(union.value_offset(1), 1);
145/// assert_eq!(union.value_offset(2), 2);
146/// ```
147#[derive(Debug, Default)]
148pub struct UnionBuilder {
149    /// The current number of slots in the array
150    len: usize,
151    /// Maps field names to `FieldData` instances which track the builders for that field
152    fields: BTreeMap<String, FieldData>,
153    /// Builder to keep track of type ids
154    type_id_builder: Int8BufferBuilder,
155    /// Builder to keep track of offsets (`None` for sparse unions)
156    value_offset_builder: Option<Int32BufferBuilder>,
157    initial_capacity: usize,
158}
159
160impl UnionBuilder {
161    /// Creates a new dense array builder.
162    pub fn new_dense() -> Self {
163        Self::with_capacity_dense(1024)
164    }
165
166    /// Creates a new sparse array builder.
167    pub fn new_sparse() -> Self {
168        Self::with_capacity_sparse(1024)
169    }
170
171    /// Creates a new dense array builder with capacity.
172    pub fn with_capacity_dense(capacity: usize) -> Self {
173        Self {
174            len: 0,
175            fields: Default::default(),
176            type_id_builder: Int8BufferBuilder::new(capacity),
177            value_offset_builder: Some(Int32BufferBuilder::new(capacity)),
178            initial_capacity: capacity,
179        }
180    }
181
182    /// Creates a new sparse array builder  with capacity.
183    pub fn with_capacity_sparse(capacity: usize) -> Self {
184        Self {
185            len: 0,
186            fields: Default::default(),
187            type_id_builder: Int8BufferBuilder::new(capacity),
188            value_offset_builder: None,
189            initial_capacity: capacity,
190        }
191    }
192
193    /// Appends a null to this builder, encoding the null in the array
194    /// of the `type_name` child / field.
195    ///
196    /// Since `UnionArray` encodes nulls as an entry in its children
197    /// (it doesn't have a validity bitmap itself), and where the null
198    /// is part of the final array, appending a NULL requires
199    /// specifying which field (child) to use.
200    #[inline]
201    pub fn append_null<T: ArrowPrimitiveType>(
202        &mut self,
203        type_name: &str,
204    ) -> Result<(), ArrowError> {
205        self.append_option::<T>(type_name, None)
206    }
207
208    /// Appends a value to this builder.
209    #[inline]
210    pub fn append<T: ArrowPrimitiveType>(
211        &mut self,
212        type_name: &str,
213        v: T::Native,
214    ) -> Result<(), ArrowError> {
215        self.append_option::<T>(type_name, Some(v))
216    }
217
218    fn append_option<T: ArrowPrimitiveType>(
219        &mut self,
220        type_name: &str,
221        v: Option<T::Native>,
222    ) -> Result<(), ArrowError> {
223        let type_name = type_name.to_string();
224
225        let mut field_data = match self.fields.remove(&type_name) {
226            Some(data) => {
227                if data.data_type != T::DATA_TYPE {
228                    return Err(ArrowError::InvalidArgumentError(format!(
229                        "Attempt to write col \"{}\" with type {} doesn't match existing type {}",
230                        type_name,
231                        T::DATA_TYPE,
232                        data.data_type
233                    )));
234                }
235                data
236            }
237            None => match self.value_offset_builder {
238                Some(_) => FieldData::new::<T>(
239                    self.fields.len() as i8,
240                    T::DATA_TYPE,
241                    self.initial_capacity,
242                ),
243                // In the case of a sparse union, we should pass the maximum of the currently length and the capacity.
244                None => {
245                    let mut fd = FieldData::new::<T>(
246                        self.fields.len() as i8,
247                        T::DATA_TYPE,
248                        self.len.max(self.initial_capacity),
249                    );
250                    for _ in 0..self.len {
251                        fd.append_null();
252                    }
253                    fd
254                }
255            },
256        };
257        self.type_id_builder.append(field_data.type_id);
258
259        match &mut self.value_offset_builder {
260            // Dense Union
261            Some(offset_builder) => {
262                offset_builder.append(field_data.slots as i32);
263            }
264            // Sparse Union
265            None => {
266                for fd in self.fields.values_mut() {
267                    // Append to all bar the FieldData currently being appended to
268                    fd.append_null();
269                }
270            }
271        }
272
273        match v {
274            Some(v) => field_data.append_value::<T>(v),
275            None => field_data.append_null(),
276        }
277
278        self.fields.insert(type_name, field_data);
279        self.len += 1;
280        Ok(())
281    }
282
283    /// Builds this builder creating a new `UnionArray`.
284    pub fn build(self) -> Result<UnionArray, ArrowError> {
285        let mut children = Vec::with_capacity(self.fields.len());
286        let union_fields = self
287            .fields
288            .into_iter()
289            .map(
290                |(
291                    name,
292                    FieldData {
293                        type_id,
294                        data_type,
295                        mut values_buffer,
296                        slots,
297                        mut null_buffer_builder,
298                    },
299                )| {
300                    // SAFETY: builder is constructed from valid value buffer, slot count, and null buffer maintained by the builder
301                    let array_ref = make_array(unsafe {
302                        ArrayDataBuilder::new(data_type.clone())
303                            .add_buffer(values_buffer.finish())
304                            .len(slots)
305                            .nulls(null_buffer_builder.finish())
306                            .build_unchecked()
307                    });
308                    children.push(array_ref);
309                    (type_id, Arc::new(Field::new(name, data_type, false)))
310                },
311            )
312            .collect();
313        UnionArray::try_new(
314            union_fields,
315            self.type_id_builder.into(),
316            self.value_offset_builder.map(Into::into),
317            children,
318        )
319    }
320
321    /// Builds this builder creating a new `UnionArray` without consuming the builder.
322    ///
323    /// This is used for the `finish_cloned` implementation in `ArrayBuilder`.
324    fn build_cloned(&self) -> Result<UnionArray, ArrowError> {
325        let mut children = Vec::with_capacity(self.fields.len());
326        let union_fields: Vec<_> = self
327            .fields
328            .iter()
329            .map(|(name, field_data)| {
330                let FieldData {
331                    type_id,
332                    data_type,
333                    values_buffer,
334                    slots,
335                    null_buffer_builder,
336                } = field_data;
337
338                // SAFETY: builder is constructed from valid value buffer, slot count, and null buffer maintained by the builder
339                let array_ref = make_array(unsafe {
340                    ArrayDataBuilder::new(data_type.clone())
341                        .add_buffer(values_buffer.finish_cloned())
342                        .len(*slots)
343                        .nulls(null_buffer_builder.finish_cloned())
344                        .build_unchecked()
345                });
346                children.push(array_ref);
347                (
348                    *type_id,
349                    Arc::new(Field::new(name.clone(), data_type.clone(), false)),
350                )
351            })
352            .collect();
353        UnionArray::try_new(
354            union_fields.into_iter().collect(),
355            ScalarBuffer::from(self.type_id_builder.as_slice().to_vec()),
356            self.value_offset_builder
357                .as_ref()
358                .map(|builder| ScalarBuffer::from(builder.as_slice().to_vec())),
359            children,
360        )
361    }
362}
363
364impl ArrayBuilder for UnionBuilder {
365    /// Returns the number of array slots in the builder
366    fn len(&self) -> usize {
367        self.len
368    }
369
370    /// Builds the array
371    fn finish(&mut self) -> ArrayRef {
372        // Even simpler - just move the builder using mem::take and replace with default
373        let builder = std::mem::take(self);
374
375        // Since UnionBuilder controls all invariants, this should never fail
376        Arc::new(builder.build().unwrap())
377    }
378
379    /// Builds the array without resetting the underlying builder
380    fn finish_cloned(&self) -> ArrayRef {
381        // We construct the UnionArray carefully to ensure try_new cannot fail.
382        // Since UnionBuilder controls all the invariants, this should never panic.
383        Arc::new(
384            self.build_cloned().unwrap_or_else(|err| {
385                panic!("UnionBuilder::build_cloned failed unexpectedly: {err}")
386            }),
387        )
388    }
389
390    /// Returns the builder as a non-mutable `Any` reference
391    fn as_any(&self) -> &dyn Any {
392        self
393    }
394
395    /// Returns the builder as a mutable `Any` reference
396    fn as_any_mut(&mut self) -> &mut dyn Any {
397        self
398    }
399
400    /// Returns the boxed builder as a box of `Any`
401    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
402        self
403    }
404}
405
406#[cfg(test)]
407mod tests {
408    use super::*;
409    use crate::array::Array;
410    use crate::cast::AsArray;
411    use crate::types::{Float64Type, Int32Type};
412
413    #[test]
414    fn test_union_builder_array_builder_trait() {
415        // Test that UnionBuilder implements ArrayBuilder trait
416        let mut builder = UnionBuilder::new_dense();
417
418        // Add some data
419        builder.append::<Int32Type>("a", 1).unwrap();
420        builder.append::<Float64Type>("b", 3.0).unwrap();
421        builder.append::<Int32Type>("a", 4).unwrap();
422
423        assert_eq!(builder.len(), 3);
424
425        // Test finish_cloned (non-destructive)
426        let array1 = builder.finish_cloned();
427        assert_eq!(array1.len(), 3);
428
429        // Verify values in cloned array
430        let union1 = array1.as_any().downcast_ref::<UnionArray>().unwrap();
431        assert_eq!(union1.type_ids(), &[0, 1, 0]);
432        assert_eq!(union1.offsets().unwrap().as_ref(), &[0, 0, 1]);
433        let int_array1 = union1.child(0).as_primitive::<Int32Type>();
434        let float_array1 = union1.child(1).as_primitive::<Float64Type>();
435        assert_eq!(int_array1.value(0), 1);
436        assert_eq!(int_array1.value(1), 4);
437        assert_eq!(float_array1.value(0), 3.0);
438
439        // Builder should still be usable after finish_cloned
440        builder.append::<Float64Type>("b", 5.0).unwrap();
441        assert_eq!(builder.len(), 4);
442
443        // Test finish (destructive)
444        let array2 = builder.finish();
445        assert_eq!(array2.len(), 4);
446
447        // Verify values in final array
448        let union2 = array2.as_any().downcast_ref::<UnionArray>().unwrap();
449        assert_eq!(union2.type_ids(), &[0, 1, 0, 1]);
450        assert_eq!(union2.offsets().unwrap().as_ref(), &[0, 0, 1, 1]);
451        let int_array2 = union2.child(0).as_primitive::<Int32Type>();
452        let float_array2 = union2.child(1).as_primitive::<Float64Type>();
453        assert_eq!(int_array2.value(0), 1);
454        assert_eq!(int_array2.value(1), 4);
455        assert_eq!(float_array2.value(0), 3.0);
456        assert_eq!(float_array2.value(1), 5.0);
457    }
458
459    #[test]
460    fn test_union_builder_type_erased() {
461        // Test type-erased usage with Box<dyn ArrayBuilder>
462        let mut builders: Vec<Box<dyn ArrayBuilder>> = vec![Box::new(UnionBuilder::new_sparse())];
463
464        // Downcast and use
465        let union_builder = builders[0]
466            .as_any_mut()
467            .downcast_mut::<UnionBuilder>()
468            .unwrap();
469        union_builder.append::<Int32Type>("x", 10).unwrap();
470        union_builder.append::<Float64Type>("y", 20.0).unwrap();
471
472        assert_eq!(builders[0].len(), 2);
473
474        let result = builders
475            .into_iter()
476            .map(|mut b| b.finish())
477            .collect::<Vec<_>>();
478        assert_eq!(result[0].len(), 2);
479
480        // Verify sparse union values
481        let union = result[0].as_any().downcast_ref::<UnionArray>().unwrap();
482        assert_eq!(union.type_ids(), &[0, 1]);
483        assert!(union.offsets().is_none()); // Sparse union has no offsets
484        let int_array = union.child(0).as_primitive::<Int32Type>();
485        let float_array = union.child(1).as_primitive::<Float64Type>();
486        assert_eq!(int_array.value(0), 10);
487        assert!(int_array.is_null(1)); // Null in sparse layout
488        assert!(float_array.is_null(0)); // Null in sparse layout
489        assert_eq!(float_array.value(1), 20.0);
490    }
491}
arrow_array/builder/union_builder.rs

arrow_array/builder/
union_builder.rs
