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    pub fn finish(&mut self) -> StructArray {
240        self.validate_content();
241        if self.fields.is_empty() {
242            return StructArray::new_empty_fields(self.len(), self.null_buffer_builder.finish());
243        }
244
245        let arrays = self.field_builders.iter_mut().map(|f| f.finish()).collect();
246        let nulls = self.null_buffer_builder.finish();
247        StructArray::new(self.fields.clone(), arrays, nulls)
248    }
249
250    /// Builds the `StructArray` without resetting the builder.
251    pub fn finish_cloned(&self) -> StructArray {
252        self.validate_content();
253
254        if self.fields.is_empty() {
255            return StructArray::new_empty_fields(
256                self.len(),
257                self.null_buffer_builder.finish_cloned(),
258            );
259        }
260
261        let arrays = self
262            .field_builders
263            .iter()
264            .map(|f| f.finish_cloned())
265            .collect();
266
267        let nulls = self.null_buffer_builder.finish_cloned();
268
269        StructArray::new(self.fields.clone(), arrays, nulls)
270    }
271
272    fn finish_preserve_values(&mut self) -> StructArray {
273        self.validate_content();
274        if self.fields.is_empty() {
275            return StructArray::new_empty_fields(self.len(), self.null_buffer_builder.finish());
276        }
277
278        let arrays = self
279            .field_builders
280            .iter_mut()
281            .map(|f| f.finish_preserve_values())
282            .collect();
283
284        let nulls = self.null_buffer_builder.finish();
285
286        StructArray::new(self.fields.clone(), arrays, nulls)
287    }
288
289    /// Constructs and validates contents in the builder to ensure that
290    /// - fields and field_builders are of equal length
291    /// - the number of items in individual field_builders are equal to self.len()
292    fn validate_content(&self) {
293        if self.fields.len() != self.field_builders.len() {
294            panic!("Number of fields is not equal to the number of field_builders.");
295        }
296        self.field_builders.iter().enumerate().for_each(|(idx, x)| {
297            if x.len() != self.len() {
298                let builder = SchemaBuilder::from(&self.fields);
299                let schema = builder.finish();
300
301                panic!("{}", format!(
302                    "StructBuilder ({}) and field_builder with index {} ({}) are of unequal lengths: ({} != {}).",
303                    schema,
304                    idx,
305                    self.fields[idx].data_type(),
306                    self.len(),
307                    x.len()
308                ));
309            }
310        });
311    }
312
313    /// Returns the current null buffer as a slice
314    pub fn validity_slice(&self) -> Option<&[u8]> {
315        self.null_buffer_builder.as_slice()
316    }
317
318    /// Returns the current null buffer allocated capacity, in bytes.
319    pub fn validity_capacity(&self) -> usize {
320        self.null_buffer_builder.allocated_size()
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use std::any::type_name;
327
328    use super::*;
329    use arrow_buffer::Buffer;
330    use arrow_data::ArrayData;
331    use arrow_schema::Field;
332
333    use crate::{array::Array, builder::tests::PreserveValuesMock, types::ArrowDictionaryKeyType};
334
335    #[test]
336    fn test_struct_array_builder() {
337        let string_builder = StringBuilder::new();
338        let int_builder = Int32Builder::new();
339
340        let fields = vec![
341            Field::new("f1", DataType::Utf8, true),
342            Field::new("f2", DataType::Int32, true),
343        ];
344        let field_builders = vec![
345            Box::new(string_builder) as Box<dyn ArrayBuilder>,
346            Box::new(int_builder) as Box<dyn ArrayBuilder>,
347        ];
348
349        let mut builder = StructBuilder::new(fields, field_builders);
350        assert_eq!(2, builder.num_fields());
351
352        let string_builder = builder
353            .field_builder::<StringBuilder>(0)
354            .expect("builder at field 0 should be string builder");
355        string_builder.append_value("joe");
356        string_builder.append_null();
357        string_builder.append_null();
358        string_builder.append_value("mark");
359        string_builder.append_nulls(2);
360        string_builder.append_value("terry");
361
362        let int_builder = builder
363            .field_builder::<Int32Builder>(1)
364            .expect("builder at field 1 should be int builder");
365        int_builder.append_value(1);
366        int_builder.append_value(2);
367        int_builder.append_null();
368        int_builder.append_value(4);
369        int_builder.append_nulls(2);
370        int_builder.append_value(3);
371
372        builder.append(true);
373        builder.append(true);
374        builder.append_null();
375        builder.append(true);
376
377        builder.append_nulls(2);
378        builder.append(true);
379
380        let struct_data = builder.finish().into_data();
381
382        assert_eq!(7, struct_data.len());
383        assert_eq!(3, struct_data.null_count());
384        assert_eq!(&[75_u8], struct_data.nulls().unwrap().validity());
385
386        let expected_string_data = ArrayData::builder(DataType::Utf8)
387            .len(7)
388            .null_bit_buffer(Some(Buffer::from(&[73_u8])))
389            .add_buffer(Buffer::from_slice_ref([0, 3, 3, 3, 7, 7, 7, 12]))
390            .add_buffer(Buffer::from_slice_ref(b"joemarkterry"))
391            .build()
392            .unwrap();
393
394        let expected_int_data = ArrayData::builder(DataType::Int32)
395            .len(7)
396            .null_bit_buffer(Some(Buffer::from_slice_ref([75_u8])))
397            .add_buffer(Buffer::from_slice_ref([1, 2, 0, 4, 4, 4, 3]))
398            .build()
399            .unwrap();
400
401        assert_eq!(expected_string_data, struct_data.child_data()[0]);
402        assert_eq!(expected_int_data, struct_data.child_data()[1]);
403
404        assert!(struct_data.is_null(4));
405        assert!(struct_data.is_null(5));
406    }
407
408    #[test]
409    fn test_struct_array_builder_finish() {
410        let int_builder = Int32Builder::new();
411        let bool_builder = BooleanBuilder::new();
412
413        let fields = vec![
414            Field::new("f1", DataType::Int32, false),
415            Field::new("f2", DataType::Boolean, false),
416        ];
417        let field_builders = vec![
418            Box::new(int_builder) as Box<dyn ArrayBuilder>,
419            Box::new(bool_builder) as Box<dyn ArrayBuilder>,
420        ];
421
422        let mut builder = StructBuilder::new(fields, field_builders);
423        builder
424            .field_builder::<Int32Builder>(0)
425            .unwrap()
426            .append_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
427        builder
428            .field_builder::<BooleanBuilder>(1)
429            .unwrap()
430            .append_slice(&[
431                false, true, false, true, false, true, false, true, false, true,
432            ]);
433
434        // Append slot values - all are valid.
435        for _ in 0..10 {
436            builder.append(true);
437        }
438
439        assert_eq!(10, builder.len());
440
441        let arr = builder.finish();
442
443        assert_eq!(10, arr.len());
444        assert_eq!(0, builder.len());
445
446        builder
447            .field_builder::<Int32Builder>(0)
448            .unwrap()
449            .append_slice(&[1, 3, 5, 7, 9]);
450        builder
451            .field_builder::<BooleanBuilder>(1)
452            .unwrap()
453            .append_slice(&[false, true, false, true, false]);
454
455        // Append slot values - all are valid.
456        for _ in 0..5 {
457            builder.append(true);
458        }
459
460        assert_eq!(5, builder.len());
461
462        let arr = builder.finish();
463
464        assert_eq!(5, arr.len());
465        assert_eq!(0, builder.len());
466    }
467
468    #[test]
469    fn test_build_fixed_size_list() {
470        const LIST_LENGTH: i32 = 4;
471        let fixed_size_list_dtype =
472            DataType::new_fixed_size_list(DataType::Int32, LIST_LENGTH, false);
473        let mut builder = make_builder(&fixed_size_list_dtype, 10);
474        let builder = builder
475            .as_any_mut()
476            .downcast_mut::<FixedSizeListBuilder<Box<dyn ArrayBuilder>>>();
477        match builder {
478            Some(builder) => {
479                assert_eq!(builder.value_length(), LIST_LENGTH);
480                assert!(
481                    builder
482                        .values()
483                        .as_any_mut()
484                        .downcast_mut::<Int32Builder>()
485                        .is_some()
486                );
487            }
488            None => panic!("expected FixedSizeListBuilder, got a different builder type"),
489        }
490    }
491
492    #[test]
493    fn test_struct_array_builder_finish_cloned() {
494        let int_builder = Int32Builder::new();
495        let bool_builder = BooleanBuilder::new();
496
497        let fields = vec![
498            Field::new("f1", DataType::Int32, false),
499            Field::new("f2", DataType::Boolean, false),
500        ];
501        let field_builders = vec![
502            Box::new(int_builder) as Box<dyn ArrayBuilder>,
503            Box::new(bool_builder) as Box<dyn ArrayBuilder>,
504        ];
505
506        let mut builder = StructBuilder::new(fields, field_builders);
507        builder
508            .field_builder::<Int32Builder>(0)
509            .unwrap()
510            .append_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
511        builder
512            .field_builder::<BooleanBuilder>(1)
513            .unwrap()
514            .append_slice(&[
515                false, true, false, true, false, true, false, true, false, true,
516            ]);
517
518        // Append slot values - all are valid.
519        for _ in 0..10 {
520            builder.append(true);
521        }
522
523        assert_eq!(10, builder.len());
524
525        let mut arr = builder.finish_cloned();
526
527        assert_eq!(10, arr.len());
528        assert_eq!(10, builder.len());
529
530        builder
531            .field_builder::<Int32Builder>(0)
532            .unwrap()
533            .append_slice(&[1, 3, 5, 7, 9]);
534        builder
535            .field_builder::<BooleanBuilder>(1)
536            .unwrap()
537            .append_slice(&[false, true, false, true, false]);
538
539        // Append slot values - all are valid.
540        for _ in 0..5 {
541            builder.append(true);
542        }
543
544        assert_eq!(15, builder.len());
545
546        arr = builder.finish();
547
548        assert_eq!(15, arr.len());
549        assert_eq!(0, builder.len());
550    }
551
552    #[test]
553    fn test_struct_array_builder_finish_preserve_values() {
554        let fields = vec![Field::new("mock", DataType::Int32, false)];
555        let field_builders = vec![Box::new(PreserveValuesMock::default()) as Box<dyn ArrayBuilder>];
556
557        let mut builder = StructBuilder::new(fields, field_builders);
558        builder
559            .field_builder::<PreserveValuesMock>(0)
560            .unwrap()
561            .inner
562            .append_value(1);
563        builder.append(true);
564
565        assert_eq!(1, builder.len());
566
567        let arr = builder.finish_preserve_values();
568
569        assert_eq!(1, arr.len());
570        assert_eq!(
571            1,
572            builder
573                .field_builder::<PreserveValuesMock>(0)
574                .unwrap()
575                .called
576        );
577    }
578
579    #[test]
580    fn test_struct_array_builder_from_schema() {
581        let mut fields = vec![
582            Field::new("f1", DataType::Float32, false),
583            Field::new("f2", DataType::Utf8, false),
584        ];
585        let sub_fields = vec![
586            Field::new("g1", DataType::Int32, false),
587            Field::new("g2", DataType::Boolean, false),
588        ];
589        let struct_type = DataType::Struct(sub_fields.into());
590        fields.push(Field::new("f3", struct_type, false));
591
592        let mut builder = StructBuilder::from_fields(fields, 5);
593        assert_eq!(3, builder.num_fields());
594        assert!(builder.field_builder::<Float32Builder>(0).is_some());
595        assert!(builder.field_builder::<StringBuilder>(1).is_some());
596        assert!(builder.field_builder::<StructBuilder>(2).is_some());
597    }
598
599    #[test]
600    fn test_datatype_properties() {
601        let fields = Fields::from(vec![
602            Field::new("f1", DataType::Decimal128(1, 2), false),
603            Field::new(
604                "f2",
605                DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())),
606                false,
607            ),
608        ]);
609        let mut builder = StructBuilder::from_fields(fields.clone(), 1);
610        builder
611            .field_builder::<Decimal128Builder>(0)
612            .unwrap()
613            .append_value(1);
614        builder
615            .field_builder::<TimestampMillisecondBuilder>(1)
616            .unwrap()
617            .append_value(1);
618        builder.append(true);
619        let array = builder.finish();
620
621        assert_eq!(array.data_type(), &DataType::Struct(fields.clone()));
622        assert_eq!(array.column(0).data_type(), fields[0].data_type());
623        assert_eq!(array.column(1).data_type(), fields[1].data_type());
624    }
625
626    #[test]
627    fn test_struct_array_builder_from_dictionary_type_int8_key() {
628        test_struct_array_builder_from_dictionary_type_inner::<Int8Type>(DataType::Int8);
629    }
630
631    #[test]
632    fn test_struct_array_builder_from_dictionary_type_int16_key() {
633        test_struct_array_builder_from_dictionary_type_inner::<Int16Type>(DataType::Int16);
634    }
635
636    #[test]
637    fn test_struct_array_builder_from_dictionary_type_int32_key() {
638        test_struct_array_builder_from_dictionary_type_inner::<Int32Type>(DataType::Int32);
639    }
640
641    #[test]
642    fn test_struct_array_builder_from_dictionary_type_int64_key() {
643        test_struct_array_builder_from_dictionary_type_inner::<Int64Type>(DataType::Int64);
644    }
645
646    fn test_struct_array_builder_from_dictionary_type_inner<K: ArrowDictionaryKeyType>(
647        key_type: DataType,
648    ) {
649        let dict_field = Field::new(
650            "f1",
651            DataType::Dictionary(Box::new(key_type), Box::new(DataType::Utf8)),
652            false,
653        );
654        let fields = vec![dict_field.clone()];
655        let expected_dtype = DataType::Struct(fields.into());
656        let cloned_dict_field = dict_field.clone();
657        let expected_child_dtype = dict_field.data_type();
658        let mut struct_builder = StructBuilder::from_fields(vec![cloned_dict_field], 5);
659        let Some(dict_builder) = struct_builder.field_builder::<StringDictionaryBuilder<K>>(0)
660        else {
661            panic!(
662                "Builder should be StringDictionaryBuilder<{}>",
663                type_name::<K>()
664            )
665        };
666        dict_builder.append_value("dict string");
667        struct_builder.append(true);
668        let array = struct_builder.finish();
669
670        assert_eq!(array.data_type(), &expected_dtype);
671        assert_eq!(array.column(0).data_type(), expected_child_dtype);
672        assert_eq!(array.column(0).len(), 1);
673    }
674
675    #[test]
676    #[should_panic(
677        expected = "Data type Dictionary(UInt64, Utf8) with key type UInt64 is not currently supported"
678    )]
679    fn test_struct_array_builder_from_schema_unsupported_type() {
680        let fields = vec![
681            Field::new("f1", DataType::UInt64, false),
682            Field::new(
683                "f2",
684                DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
685                false,
686            ),
687        ];
688
689        let _ = StructBuilder::from_fields(fields, 5);
690    }
691
692    #[test]
693    #[should_panic(expected = "Dictionary value type Int32 is not currently supported")]
694    fn test_struct_array_builder_from_dict_with_unsupported_value_type() {
695        let fields = vec![Field::new(
696            "f1",
697            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)),
698            false,
699        )];
700
701        let _ = StructBuilder::from_fields(fields, 5);
702    }
703
704    #[test]
705    fn test_struct_array_builder_field_builder_type_mismatch() {
706        let int_builder = Int32Builder::with_capacity(10);
707
708        let fields = vec![Field::new("f1", DataType::Int32, false)];
709        let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
710
711        let mut builder = StructBuilder::new(fields, field_builders);
712        assert!(builder.field_builder::<BinaryBuilder>(0).is_none());
713    }
714
715    #[test]
716    #[should_panic(
717        expected = "StructBuilder (Field { \"f1\": Int32 }, Field { \"f2\": Boolean }) and field_builder with index 1 (Boolean) are of unequal lengths: (2 != 1)."
718    )]
719    fn test_struct_array_builder_unequal_field_builders_lengths() {
720        let mut int_builder = Int32Builder::with_capacity(10);
721        let mut bool_builder = BooleanBuilder::new();
722
723        int_builder.append_value(1);
724        int_builder.append_value(2);
725        bool_builder.append_value(true);
726
727        let fields = vec![
728            Field::new("f1", DataType::Int32, false),
729            Field::new("f2", DataType::Boolean, false),
730        ];
731        let field_builders = vec![
732            Box::new(int_builder) as Box<dyn ArrayBuilder>,
733            Box::new(bool_builder) as Box<dyn ArrayBuilder>,
734        ];
735
736        let mut builder = StructBuilder::new(fields, field_builders);
737        builder.append(true);
738        builder.append(true);
739        builder.finish();
740    }
741
742    #[test]
743    #[should_panic(expected = "Number of fields is not equal to the number of field_builders.")]
744    fn test_struct_array_builder_unequal_field_field_builders() {
745        let int_builder = Int32Builder::with_capacity(10);
746
747        let fields = vec![
748            Field::new("f1", DataType::Int32, false),
749            Field::new("f2", DataType::Boolean, false),
750        ];
751        let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
752
753        let mut builder = StructBuilder::new(fields, field_builders);
754        builder.finish();
755    }
756
757    #[test]
758    #[should_panic(
759        expected = "Incorrect datatype for StructArray field \\\"timestamp\\\", expected Timestamp(ns, \\\"UTC\\\") got Timestamp(ns)"
760    )]
761    fn test_struct_array_mismatch_builder() {
762        let fields = vec![Field::new(
763            "timestamp",
764            DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".to_owned().into())),
765            false,
766        )];
767
768        let field_builders: Vec<Box<dyn ArrayBuilder>> =
769            vec![Box::new(TimestampNanosecondBuilder::new())];
770
771        let mut sa = StructBuilder::new(fields, field_builders);
772        sa.finish();
773    }
774
775    #[test]
776    fn test_empty() {
777        let mut builder = StructBuilder::new(Fields::empty(), vec![]);
778        builder.append(true);
779        builder.append(false);
780
781        let a1 = builder.finish_cloned();
782        let a2 = builder.finish();
783        assert_eq!(a1, a2);
784        assert_eq!(a1.len(), 2);
785        assert_eq!(a1.null_count(), 1);
786        assert!(a1.is_valid(0));
787        assert!(a1.is_null(1));
788    }
789
790    #[test]
791    fn test_append_non_nulls() {
792        let int_builder = Int32Builder::new();
793        let fields = vec![Field::new("f1", DataType::Int32, false)];
794        let field_builders = vec![Box::new(int_builder) as Box<dyn ArrayBuilder>];
795
796        let mut builder = StructBuilder::new(fields, field_builders);
797        builder
798            .field_builder::<Int32Builder>(0)
799            .unwrap()
800            .append_slice(&[1, 2, 3, 4, 5]);
801        builder.append_non_nulls(5);
802
803        let arr = builder.finish();
804        assert_eq!(arr.len(), 5);
805        assert_eq!(arr.null_count(), 0);
806        for i in 0..5 {
807            assert!(arr.is_valid(i));
808        }
809    }
810
811    #[test]
812    fn test_append_non_nulls_with_nulls() {
813        let mut builder = StructBuilder::new(Fields::empty(), vec![]);
814        builder.append_null();
815        builder.append_non_nulls(3);
816        builder.append_nulls(2);
817        builder.append_non_nulls(1);
818
819        let arr = builder.finish();
820        assert_eq!(arr.len(), 7);
821        assert_eq!(arr.null_count(), 3);
822        assert!(arr.is_null(0));
823        assert!(arr.is_valid(1));
824        assert!(arr.is_valid(2));
825        assert!(arr.is_valid(3));
826        assert!(arr.is_null(4));
827        assert!(arr.is_null(5));
828        assert!(arr.is_valid(6));
829    }
830
831    #[test]
832    fn test_append_non_nulls_zero() {
833        let mut builder = StructBuilder::new(Fields::empty(), vec![]);
834        builder.append_non_nulls(0);
835        assert_eq!(builder.len(), 0);
836
837        builder.append(true);
838        builder.append_non_nulls(0);
839        assert_eq!(builder.len(), 1);
840
841        let arr = builder.finish();
842        assert_eq!(arr.len(), 1);
843        assert_eq!(arr.null_count(), 0);
844    }
845}
arrow_array/builder/struct_builder.rs

arrow_array/builder/
struct_builder.rs
