arrow_array/builder/
union_builder.rs

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

arrow_array/builder/
union_builder.rs
