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

arrow_array/builder/
struct_builder.rs
