arrow_array/builder/
struct_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::StructArray;
19use crate::builder::*;
20use arrow_buffer::NullBufferBuilder;
21use arrow_schema::{Fields, SchemaBuilder};
22use std::sync::Arc;
23
24/// Builder for [`StructArray`]
25///
26/// Note that callers should make sure that methods of all the child field builders are
27/// properly called to maintain the consistency of the data structure.
28///
29///
30/// Handling arrays with complex layouts, such as `List<Struct<List<Struct>>>`, in Rust can be challenging due to its strong typing system.
31/// To construct a collection builder ([`ListBuilder`], [`LargeListBuilder`], or [`MapBuilder`]) using [`make_builder`], multiple calls are required. This complexity arises from the recursive approach utilized by [`StructBuilder::from_fields`].
32///
33/// Initially, [`StructBuilder::from_fields`] invokes [`make_builder`], which returns a `Box<dyn ArrayBuilder>`. To obtain the specific collection builder, one must first use [`StructBuilder::field_builder`] to get a `Collection<[Box<dyn ArrayBuilder>]>`. Subsequently, the `values()` result from this operation can be downcast to the desired builder type.
34///
35/// For example, when working with [`ListBuilder`], you would first call [`StructBuilder::field_builder::<ListBuilder<Box<dyn ArrayBuilder>>>`] and then downcast the [`Box<dyn ArrayBuilder>`] to the specific [`StructBuilder`] you need.
36///
37/// For a practical example see the code below:
38///
39/// ```rust
40///    use arrow_array::builder::{ArrayBuilder, ListBuilder, StringBuilder, StructBuilder};
41///    use arrow_schema::{DataType, Field, Fields};
42///    use std::sync::Arc;
43///
44///    // This is an example column that has a List<Struct<List<Struct>>> layout
45///    let mut example_col = ListBuilder::new(StructBuilder::from_fields(
46///        vec![Field::new(
47///            "value_list",
48///            DataType::List(Arc::new(Field::new_list_field(
49///                DataType::Struct(Fields::from(vec![
50///                    Field::new("key", DataType::Utf8, true),
51///                    Field::new("value", DataType::Utf8, true),
52///                ])), //In this example we are trying to get to this builder and insert key/value pairs
53///                true,
54///            ))),
55///            true,
56///        )],
57///        0,
58///    ));
59///
60///   // We can obtain the StructBuilder without issues, because example_col was created with StructBuilder
61///   let col_struct_builder: &mut StructBuilder = example_col.values();
62///
63///   // We can't obtain the ListBuilder<StructBuilder> with the expected generic types, because under the hood
64///   // the StructBuilder was returned as a Box<dyn ArrayBuilder> and passed as such to the ListBuilder constructor
65///
66///   // This panics in runtime, even though we know that the builder is a ListBuilder<StructBuilder>.
67///   // let sb = col_struct_builder
68///   //     .field_builder::<ListBuilder<StructBuilder>>(0)
69///   //     .as_mut()
70///   //     .unwrap();
71///
72///   //To keep in line with Rust's strong typing, we fetch a ListBuilder<Box<dyn ArrayBuilder>> from the column StructBuilder first...
73///   let mut list_builder_option =
74///       col_struct_builder.field_builder::<ListBuilder<Box<dyn ArrayBuilder>>>(0);
75///
76///   let list_builder = list_builder_option.as_mut().unwrap();
77///
78///   // ... and then downcast the key/value pair values to a StructBuilder
79///   let struct_builder = list_builder
80///       .values()
81///       .as_any_mut()
82///       .downcast_mut::<StructBuilder>()
83///       .unwrap();
84///
85///   // We can now append values to the StructBuilder
86///   let key_builder = struct_builder.field_builder::<StringBuilder>(0).unwrap();
87///   key_builder.append_value("my key");
88///
89///   let value_builder = struct_builder.field_builder::<StringBuilder>(1).unwrap();
90///   value_builder.append_value("my value");
91///
92///   struct_builder.append(true);
93///   list_builder.append(true);
94///   col_struct_builder.append(true);
95///   example_col.append(true);
96///
97///   let array = example_col.finish();
98///
99///   println!("My array: {:?}", array);
100/// ```
101///
102pub struct StructBuilder {
103    fields: Fields,
104    field_builders: Vec<Box<dyn ArrayBuilder>>,
105    null_buffer_builder: NullBufferBuilder,
106}
107
108impl std::fmt::Debug for StructBuilder {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        f.debug_struct("StructBuilder")
111            .field("fields", &self.fields)
112            .field("bitmap_builder", &self.null_buffer_builder)
113            .field("len", &self.len())
114            .finish()
115    }
116}
117
118impl ArrayBuilder for StructBuilder {
119    /// Returns the number of array slots in the builder.
120    ///
121    /// Note that this always return the first child field builder's length, and it is
122    /// the caller's responsibility to maintain the consistency that all the child field
123    /// builder should have the equal number of elements.
124    fn len(&self) -> usize {
125        self.null_buffer_builder.len()
126    }
127
128    /// Builds the array.
129    fn finish(&mut self) -> ArrayRef {
130        Arc::new(self.finish())
131    }
132
133    /// Builds the array without resetting the builder.
134    fn finish_cloned(&self) -> ArrayRef {
135        Arc::new(self.finish_cloned())
136    }
137
138    fn finish_preserve_values(&mut self) -> ArrayRef {
139        Arc::new(self.finish_preserve_values())
140    }
141
142    /// Returns the builder as a non-mutable `Any` reference.
143    ///
144    /// This is most useful when one wants to call non-mutable APIs on a specific builder
145    /// type. In this case, one can first cast this into a `Any`, and then use
146    /// `downcast_ref` to get a reference on the specific builder.
147    fn as_any(&self) -> &dyn Any {
148        self
149    }
150
151    /// Returns the builder as a mutable `Any` reference.
152    ///
153    /// This is most useful when one wants to call mutable APIs on a specific builder
154    /// type. In this case, one can first cast this into a `Any`, and then use
155    /// `downcast_mut` to get a reference on the specific builder.
156    fn as_any_mut(&mut self) -> &mut dyn Any {
157        self
158    }
159
160    /// Returns the boxed builder as a box of `Any`.
161    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
162        self
163    }
164}
165
166impl StructBuilder {
167    /// Creates a new `StructBuilder`
168    pub fn new(fields: impl Into<Fields>, field_builders: Vec<Box<dyn ArrayBuilder>>) -> Self {
169        Self {
170            field_builders,
171            fields: fields.into(),
172            null_buffer_builder: NullBufferBuilder::new(0),
173        }
174    }
175
176    /// Creates a new `StructBuilder` from [`Fields`] and `capacity`
177    pub fn from_fields(fields: impl Into<Fields>, capacity: usize) -> Self {
178        let fields = fields.into();
179        let mut builders = Vec::with_capacity(fields.len());
180        for field in &fields {
181            builders.push(make_builder(field.data_type(), capacity));
182        }
183        Self::new(fields, builders)
184    }
185
186    /// Returns a mutable reference to the child field builder at index `i`.
187    /// Result will be `None` if the input type `T` provided doesn't match the actual
188    /// field builder's type.
189    pub fn field_builder<T: ArrayBuilder>(&mut self, i: usize) -> Option<&mut T> {
190        self.field_builders[i].as_any_mut().downcast_mut::<T>()
191    }
192
193    /// Returns a reference to field builders
194    pub fn field_builders(&self) -> &[Box<dyn ArrayBuilder>] {
195        &self.field_builders
196    }
197
198    /// Returns a mutable reference to field builders
199    pub fn field_builders_mut(&mut self) -> &mut [Box<dyn ArrayBuilder>] {
200        &mut self.field_builders
201    }
202
203    /// Returns the number of fields for the struct this builder is building.
204    pub fn num_fields(&self) -> usize {
205        self.field_builders.len()
206    }
207
208    /// Returns the fields for the struct this builder is building.
209    pub fn fields(&self) -> &Fields {
210        &self.fields
211    }
212
213    /// Appends an element (either null or non-null) to the struct. The actual elements
214    /// should be appended for each child sub-array in a consistent way.
215    #[inline]
216    pub fn append(&mut self, is_valid: bool) {
217        self.null_buffer_builder.append(is_valid);
218    }
219
220    /// Appends `n` non-null entries into the builder.
221    #[inline]
222    pub fn append_non_nulls(&mut self, n: usize) {
223        self.null_buffer_builder.append_n_non_nulls(n);
224    }
225
226    /// Appends a null element to the struct.
227    #[inline]
228    pub fn append_null(&mut self) {
229        self.append(false)
230    }
231
232    /// Appends `n` `null`s into the builder.
233    #[inline]
234    pub fn append_nulls(&mut self, n: usize) {
235        self.null_buffer_builder.append_n_nulls(n);
236    }
237
238    /// Builds the `StructArray` and reset this builder.
239    ///
240    /// # Panics
241    ///
242    /// Panics if the number of fields is not equal to the number of field builders, or
243    /// if the field builders do not all have length `self.len()`
244    pub fn finish(&mut self) -> StructArray {
245        self.validate_content();
246        if self.fields.is_empty() {
247            return StructArray::new_empty_fields(self.len(), self.null_buffer_builder.finish());
248        }
249
250        let arrays = self.field_builders.iter_mut().map(|f| f.finish()).collect();
251        let nulls = self.null_buffer_builder.finish();
252        StructArray::new(self.fields.clone(), arrays, nulls)
253    }
254
255    /// Builds the `StructArray` without resetting the builder.
256    ///
257    /// # Panics
258    ///
259    /// Panics if the number of fields is not equal to the number of field builders, or
260    /// if the field builders do not all have length `self.len()`
261    pub fn finish_cloned(&self) -> StructArray {
262        self.validate_content();
263
264        if self.fields.is_empty() {
265            return StructArray::new_empty_fields(
266                self.len(),
267                self.null_buffer_builder.finish_cloned(),
268            );
269        }
270
271        let arrays = self
272            .field_builders
273            .iter()
274            .map(|f| f.finish_cloned())
275            .collect();
276
277        let nulls = self.null_buffer_builder.finish_cloned();
278
279        StructArray::new(self.fields.clone(), arrays, nulls)
280    }
281
282    fn finish_preserve_values(&mut self) -> StructArray {
283        self.validate_content();
284        if self.fields.is_empty() {
285            return StructArray::new_empty_fields(self.len(), self.null_buffer_builder.finish());
286        }
287
288        let arrays = self
289            .field_builders
290            .iter_mut()
291            .map(|f| f.finish_preserve_values())
292            .collect();
293
294        let nulls = self.null_buffer_builder.finish();
295
296        StructArray::new(self.fields.clone(), arrays, nulls)
297    }
298
299    /// Constructs and validates contents in the builder to ensure that
300    /// - fields and field_builders are of equal length
301    /// - the number of items in individual field_builders are equal to self.len()
302    fn validate_content(&self) {
303        if self.fields.len() != self.field_builders.len() {
304            panic!("Number of fields is not equal to the number of field_builders.");
305        }
306        self.field_builders.iter().enumerate().for_each(|(idx, x)| {
307            if x.len() != self.len() {
308                let builder = SchemaBuilder::from(&self.fields);
309                let schema = builder.finish();
310
311                panic!("{}", format!(
312                    "StructBuilder ({}) and field_builder with index {} ({}) are of unequal lengths: ({} != {}).",
313                    schema,
314                    idx,
315                    self.fields[idx].data_type(),
316                    self.len(),
317                    x.len()
318                ));
319            }
320        });
321    }
322
323    /// Returns the current null buffer as a slice
324    pub fn validity_slice(&self) -> Option<&[u8]> {
325        self.null_buffer_builder.as_slice()
326    }
327
328    /// Returns the current null buffer allocated capacity, in bytes.
329    pub fn validity_capacity(&self) -> usize {
330        self.null_buffer_builder.allocated_size()
331    }
332}
333
334#[cfg(test)]
335mod tests {
336    use std::any::type_name;
337
338    use super::*;
339    use arrow_buffer::Buffer;
340    use arrow_data::ArrayData;
341    use arrow_schema::Field;
342
343    use crate::{array::Array, builder::tests::PreserveValuesMock, types::ArrowDictionaryKeyType};
344
345    #[test]
346    fn test_struct_array_builder() {
347        let string_builder = StringBuilder::new();
348        let int_builder = Int32Builder::new();
349
350        let fields = vec![
351            Field::new("f1", DataType::Utf8, true),
352            Field::new("f2", DataType::Int32, true),
353        ];
354        let field_builders = vec![
355            Box::new(string_builder) as Box<dyn ArrayBuilder>,
356            Box::new(int_builder) as Box<dyn ArrayBuilder>,
357        ];
358
359        let mut builder = StructBuilder::new(fields, field_builders);
360        assert_eq!(2, builder.num_fields());
361
362        let string_builder = builder
363            .field_builder::<StringBuilder>(0)
364            .expect("builder at field 0 should be string builder");
365        string_builder.append_value("joe");
366        string_builder.append_null();
367        string_builder.append_null();
368        string_builder.append_value("mark");
369        string_builder.append_nulls(2);
370        string_builder.append_value("terry");
371
372        let int_builder = builder
373            .field_builder::<Int32Builder>(1)
374            .expect("builder at field 1 should be int builder");
375        int_builder.append_value(1);
376        int_builder.append_value(2);
377        int_builder.append_null();
378        int_builder.append_value(4);
379        int_builder.append_nulls(2);
380        int_builder.append_value(3);
381
382        builder.append(true);
383        builder.append(true);
384        builder.append_null();
385        builder.append(true);
386
387        builder.append_nulls(2);
388        builder.append(true);
389
390        let struct_data = builder.finish().into_data();
391
392        assert_eq!(7, struct_data.len());
393        assert_eq!(3, struct_data.null_count());
394        assert_eq!(&[75_u8], struct_data.nulls().unwrap().validity());
395
396        let expected_string_data = ArrayData::builder(DataType::Utf8)
397            .len(7)
398            .null_bit_buffer(Some(Buffer::from(&[73_u8])))
399            .add_buffer(Buffer::from_slice_ref([0, 3, 3, 3, 7, 7, 7, 12]))
400            .add_buffer(Buffer::from_slice_ref(b"joemarkterry"))
401            .build()
402            .unwrap();
403
404        let expected_int_data = ArrayData::builder(DataType::Int32)
405            .len(7)
406            .null_bit_buffer(Some(Buffer::from_slice_ref([75_u8])))
407            .add_buffer(Buffer::from_slice_ref([1, 2, 0, 4, 4, 4, 3]))
408            .build()
409            .unwrap();
410
411        assert_eq!(expected_string_data, struct_data.child_data()[0]);
412        assert_eq!(expected_int_data, struct_data.child_data()[1]);
413
414        assert!(struct_data.is_null(4));
415        assert!(struct_data.is_null(5));
416    }
417
418    #[test]
419    fn test_struct_array_builder_finish() {
420        let int_builder = Int32Builder::new();
421        let bool_builder = BooleanBuilder::new();
422
423        let fields = vec![
424            Field::new("f1", DataType::Int32, false),
425            Field::new("f2", DataType::Boolean, false),
426        ];
427        let field_builders = vec![
428            Box::new(int_builder) as Box<dyn ArrayBuilder>,
429            Box::new(bool_builder) as Box<dyn ArrayBuilder>,
430        ];
431
432        let mut builder = StructBuilder::new(fields, field_builders);
433        builder
434            .field_builder::<Int32Builder>(0)
435            .unwrap()
436            .append_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
437        builder
438            .field_builder::<BooleanBuilder>(1)
439            .unwrap()
440            .append_slice(&[
441                false, true, false, true, false, true, false, true, false, true,
442            ]);
443
444        // Append slot values - all are valid.
445        for _ in 0..10 {
446            builder.append(true);
447        }
448
449        assert_eq!(10, builder.len());
450
451        let arr = builder.finish();
452
453        assert_eq!(10, arr.len());
454        assert_eq!(0, builder.len());
455
456        builder
457            .field_builder::<Int32Builder>(0)
458            .unwrap()
459            .append_slice(&[1, 3, 5, 7, 9]);
460        builder
461            .field_builder::<BooleanBuilder>(1)
462            .unwrap()
463            .append_slice(&[false, true, false, true, false]);
464
465        // Append slot values - all are valid.
466        for _ in 0..5 {
467            builder.append(true);
468        }
469
470        assert_eq!(5, builder.len());
471
472        let arr = builder.finish();
473
474        assert_eq!(5, arr.len());
475        assert_eq!(0, builder.len());
476    }
477
478    #[test]
479    fn test_build_fixed_size_list() {
480        const LIST_LENGTH: i32 = 4;
481        let fixed_size_list_dtype =
482            DataType::new_fixed_size_list(DataType::Int32, LIST_LENGTH, false);
483        let mut builder = make_builder(&fixed_size_list_dtype, 10);
484        let builder = builder
485            .as_any_mut()
486            .downcast_mut::<FixedSizeListBuilder<Box<dyn ArrayBuilder>>>();
487        match builder {
488            Some(builder) => {
489                assert_eq!(builder.value_length(), LIST_LENGTH);
490                assert!(
491                    builder
492                        .values()
493                        .as_any_mut()
494                        .downcast_mut::<Int32Builder>()
495                        .is_some()
496                );
497            }
498            None => panic!("expected FixedSizeListBuilder, got a different builder type"),
499        }
500    }
501
502    #[test]
503    fn test_struct_array_builder_finish_cloned() {
504        let int_builder = Int32Builder::new();
505        let bool_builder = BooleanBuilder::new();
506
507        let fields = vec![
508            Field::new("f1", DataType::Int32, false),
509            Field::new("f2", DataType::Boolean, false),
510        ];
511        let field_builders = vec![
512            Box::new(int_builder) as Box<dyn ArrayBuilder>,
513            Box::new(bool_builder) as Box<dyn ArrayBuilder>,
514        ];
515
516        let mut builder = StructBuilder::new(fields, field_builders);
517        builder
518            .field_builder::<Int32Builder>(0)
519            .unwrap()
520            .append_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
521        builder
522            .field_builder::<BooleanBuilder>(1)
523            .unwrap()
524            .append_slice(&[
525                false, true, false, true, false, true, false, true, false, true,
526            ]);
527
528        // Append slot values - all are valid.
529        for _ in 0..10 {
530            builder.append(true);
531        }
532
533        assert_eq!(10, builder.len());
534
535        let mut arr = builder.finish_cloned();
536
537        assert_eq!(10, arr.len());
538        assert_eq!(10, builder.len());
539
540        builder
541            .field_builder::<Int32Builder>(0)
542            .unwrap()
543            .append_slice(&[1, 3, 5, 7, 9]);
544        builder
545            .field_builder::<BooleanBuilder>(1)
546            .unwrap()
547            .append_slice(&[false, true, false, true, false]);
548
549        // Append slot values - all are valid.
550        for _ in 0..5 {
551            builder.append(true);
552        }
553
554        assert_eq!(15, builder.len());
555
556        arr = builder.finish();
557
558        assert_eq!(15, arr.len());
559        assert_eq!(0, builder.len());
560    }
561
562    #[test]
563    fn test_struct_array_builder_finish_preserve_values() {
564        let fields = vec![Field::new("mock", DataType::Int32, false)];
565        let field_builders = vec![Box::new(PreserveValuesMock::default()) as Box<dyn ArrayBuilder>];
566
567        let mut builder = StructBuilder::new(fields, field_builders);
568        builder
569            .field_builder::<PreserveValuesMock>(0)
570            .unwrap()
571            .inner
572            .append_value(1);
573        builder.append(true);
574
575        assert_eq!(1, builder.len());
576
577        let arr = builder.finish_preserve_values();
578
579        assert_eq!(1, arr.len());
580        assert_eq!(
581            1,
582            builder
583                .field_builder::<PreserveValuesMock>(0)
584                .unwrap()
585                .called
586        );
587    }
588
589    #[test]
590    fn test_struct_array_builder_from_schema() {
591        let mut fields = vec![
592            Field::new("f1", DataType::Float32, false),
593            Field::new("f2", DataType::Utf8, false),
594        ];
595        let sub_fields = vec![
596            Field::new("g1", DataType::Int32, false),
597            Field::new("g2", DataType::Boolean, false),
598        ];
599        let struct_type = DataType::Struct(sub_fields.into());
600        fields.push(Field::new("f3", struct_type, false));
601
602        let mut builder = StructBuilder::from_fields(fields, 5);
603        assert_eq!(3, builder.num_fields());
604        assert!(builder.field_builder::<Float32Builder>(0).is_some());
605        assert!(builder.field_builder::<StringBuilder>(1).is_some());
606        assert!(builder.field_builder::<StructBuilder>(2).is_some());
607    }
608
609    #[test]
610    fn test_datatype_properties() {
611        let fields = Fields::from(vec![
612            Field::new("f1", DataType::Decimal128(1, 2), false),
613            Field::new(
614                "f2",
615                DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())),
616                false,
617            ),
618        ]);
619        let mut builder = StructBuilder::from_fields(fields.clone(), 1);
620        builder
621            .field_builder::<Decimal128Builder>(0)
622            .unwrap()
623            .append_value(1);
624        builder
625            .field_builder::<TimestampMillisecondBuilder>(1)
626            .unwrap()
627            .append_value(1);
628        builder.append(true);
629        let array = builder.finish();
630
631        assert_eq!(array.data_type(), &DataType::Struct(fields.clone()));
632        assert_eq!(array.column(0).data_type(), fields[0].data_type());
633        assert_eq!(array.column(1).data_type(), fields[1].data_type());
634    }
635
636    #[test]
637    fn test_struct_array_builder_from_dictionary_type_int8_key() {
638        test_struct_array_builder_from_dictionary_type_inner::<Int8Type>(DataType::Int8);
639    }
640
641    #[test]
642    fn test_struct_array_builder_from_dictionary_type_int16_key() {
643        test_struct_array_builder_from_dictionary_type_inner::<Int16Type>(DataType::Int16);
644    }
645
646    #[test]
647    fn test_struct_array_builder_from_dictionary_type_int32_key() {
648        test_struct_array_builder_from_dictionary_type_inner::<Int32Type>(DataType::Int32);
649    }
650
651    #[test]
652    fn test_struct_array_builder_from_dictionary_type_int64_key() {
653        test_struct_array_builder_from_dictionary_type_inner::<Int64Type>(DataType::Int64);
654    }
655
656    fn test_struct_array_builder_from_dictionary_type_inner<K: ArrowDictionaryKeyType>(
657        key_type: DataType,
658    ) {
659        let dict_field = Field::new(
660            "f1",
661            DataType::Dictionary(Box::new(key_type), Box::new(DataType::Utf8)),
662            false,
663        );
664        let fields = vec![dict_field.clone()];
665        let expected_dtype = DataType::Struct(fields.into());
666        let cloned_dict_field = dict_field.clone();
667        let expected_child_dtype = dict_field.data_type();
668        let mut struct_builder = StructBuilder::from_fields(vec![cloned_dict_field], 5);
669        let Some(dict_builder) = struct_builder.field_builder::<StringDictionaryBuilder<K>>(0)
670        else {
671            panic!(
672                "Builder should be StringDictionaryBuilder<{}>",
673                type_name::<K>()
674            )
675        };
676        dict_builder.append_value("dict string");
677        struct_builder.append(true);
678        let array = struct_builder.finish();
679
680        assert_eq!(array.data_type(), &expected_dtype);
681        assert_eq!(array.column(0).data_type(), expected_child_dtype);
682        assert_eq!(array.column(0).len(), 1);
683    }
684
685    #[test]
686    #[should_panic(
687        expected = "Data type Dictionary(UInt64, Utf8) with key type UInt64 is not currently supported"
688    )]
689    fn test_struct_array_builder_from_schema_unsupported_type() {
690        let fields = vec![
691            Field::new("f1", DataType::UInt64, false),
692            Field::new(
693                "f2",
694                DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
695                false,
696            ),
697        ];
698
699        let _ = StructBuilder::from_fields(fields, 5);
700    }
701
702    #[test]
703    #[should_panic(expected = "Dictionary value type Int32 is not currently supported")]
704    fn test_struct_array_builder_from_dict_with_unsupported_value_type() {
705        let fields = vec![Field::new(
706            "f1",
707            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)),
708            false,
709        )];
710
711        let _ = StructBuilder::from_fields(fields, 5);
712    }
713
714    #[test]
715    fn test_struct_array_builder_field_builder_type_mismatch() {
716        let int_builder = Int32Builder::with_capacity(10);
717
718        let fields = vec![Field::new("f1", DataType::Int32, false)];
719        let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
720
721        let mut builder = StructBuilder::new(fields, field_builders);
722        assert!(builder.field_builder::<BinaryBuilder>(0).is_none());
723    }
724
725    #[test]
726    #[should_panic(
727        expected = "StructBuilder (Field { \"f1\": Int32 }, Field { \"f2\": Boolean }) and field_builder with index 1 (Boolean) are of unequal lengths: (2 != 1)."
728    )]
729    fn test_struct_array_builder_unequal_field_builders_lengths() {
730        let mut int_builder = Int32Builder::with_capacity(10);
731        let mut bool_builder = BooleanBuilder::new();
732
733        int_builder.append_value(1);
734        int_builder.append_value(2);
735        bool_builder.append_value(true);
736
737        let fields = vec![
738            Field::new("f1", DataType::Int32, false),
739            Field::new("f2", DataType::Boolean, false),
740        ];
741        let field_builders = vec![
742            Box::new(int_builder) as Box<dyn ArrayBuilder>,
743            Box::new(bool_builder) as Box<dyn ArrayBuilder>,
744        ];
745
746        let mut builder = StructBuilder::new(fields, field_builders);
747        builder.append(true);
748        builder.append(true);
749        builder.finish();
750    }
751
752    #[test]
753    #[should_panic(expected = "Number of fields is not equal to the number of field_builders.")]
754    fn test_struct_array_builder_unequal_field_field_builders() {
755        let int_builder = Int32Builder::with_capacity(10);
756
757        let fields = vec![
758            Field::new("f1", DataType::Int32, false),
759            Field::new("f2", DataType::Boolean, false),
760        ];
761        let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
762
763        let mut builder = StructBuilder::new(fields, field_builders);
764        builder.finish();
765    }
766
767    #[test]
768    #[should_panic(
769        expected = "Incorrect datatype for StructArray field \\\"timestamp\\\", expected Timestamp(ns, \\\"UTC\\\") got Timestamp(ns)"
770    )]
771    fn test_struct_array_mismatch_builder() {
772        let fields = vec![Field::new(
773            "timestamp",
774            DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".to_owned().into())),
775            false,
776        )];
777
778        let field_builders: Vec<Box<dyn ArrayBuilder>> =
779            vec![Box::new(TimestampNanosecondBuilder::new())];
780
781        let mut sa = StructBuilder::new(fields, field_builders);
782        sa.finish();
783    }
784
785    #[test]
786    fn test_empty() {
787        let mut builder = StructBuilder::new(Fields::empty(), vec![]);
788        builder.append(true);
789        builder.append(false);
790
791        let a1 = builder.finish_cloned();
792        let a2 = builder.finish();
793        assert_eq!(a1, a2);
794        assert_eq!(a1.len(), 2);
795        assert_eq!(a1.null_count(), 1);
796        assert!(a1.is_valid(0));
797        assert!(a1.is_null(1));
798    }
799
800    #[test]
801    fn test_append_non_nulls() {
802        let int_builder = Int32Builder::new();
803        let fields = vec![Field::new("f1", DataType::Int32, false)];
804        let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
805
806        let mut builder = StructBuilder::new(fields, field_builders);
807        builder
808            .field_builder::<Int32Builder>(0)
809            .unwrap()
810            .append_slice(&[1, 2, 3, 4, 5]);
811        builder.append_non_nulls(5);
812
813        let arr = builder.finish();
814        assert_eq!(arr.len(), 5);
815        assert_eq!(arr.null_count(), 0);
816        for i in 0..5 {
817            assert!(arr.is_valid(i));
818        }
819    }
820
821    #[test]
822    fn test_append_non_nulls_with_nulls() {
823        let mut builder = StructBuilder::new(Fields::empty(), vec![]);
824        builder.append_null();
825        builder.append_non_nulls(3);
826        builder.append_nulls(2);
827        builder.append_non_nulls(1);
828
829        let arr = builder.finish();
830        assert_eq!(arr.len(), 7);
831        assert_eq!(arr.null_count(), 3);
832        assert!(arr.is_null(0));
833        assert!(arr.is_valid(1));
834        assert!(arr.is_valid(2));
835        assert!(arr.is_valid(3));
836        assert!(arr.is_null(4));
837        assert!(arr.is_null(5));
838        assert!(arr.is_valid(6));
839    }
840
841    #[test]
842    fn test_append_non_nulls_zero() {
843        let mut builder = StructBuilder::new(Fields::empty(), vec![]);
844        builder.append_non_nulls(0);
845        assert_eq!(builder.len(), 0);
846
847        builder.append(true);
848        builder.append_non_nulls(0);
849        assert_eq!(builder.len(), 1);
850
851        let arr = builder.finish();
852        assert_eq!(arr.len(), 1);
853        assert_eq!(arr.null_count(), 0);
854    }
855}
arrow_array/builder/struct_builder.rs

arrow_array/builder/
struct_builder.rs
