arrow_array/builder/
primitive_builder.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::builder::ArrayBuilder;
6use crate::types::*;
7use crate::{Array, ArrayRef, PrimitiveArray};
8use arrow_buffer::{Buffer, MutableBuffer, NullBufferBuilder, ScalarBuffer};
9use arrow_data::ArrayData;
10use arrow_schema::{ArrowError, DataType};
11use std::any::Any;
12use std::sync::Arc;
13
14/// A signed 8-bit integer array builder.
15pub type Int8Builder = PrimitiveBuilder<Int8Type>;
16/// A signed 16-bit integer array builder.
17pub type Int16Builder = PrimitiveBuilder<Int16Type>;
18/// A signed 32-bit integer array builder.
19pub type Int32Builder = PrimitiveBuilder<Int32Type>;
20/// A signed 64-bit integer array builder.
21pub type Int64Builder = PrimitiveBuilder<Int64Type>;
22/// An unsigned 8-bit integer array builder.
23pub type UInt8Builder = PrimitiveBuilder<UInt8Type>;
24/// An unsigned 16-bit integer array builder.
25pub type UInt16Builder = PrimitiveBuilder<UInt16Type>;
26/// An unsigned 32-bit integer array builder.
27pub type UInt32Builder = PrimitiveBuilder<UInt32Type>;
28/// An unsigned 64-bit integer array builder.
29pub type UInt64Builder = PrimitiveBuilder<UInt64Type>;
30/// A 16-bit floating point array builder.
31pub type Float16Builder = PrimitiveBuilder<Float16Type>;
32/// A 32-bit floating point array builder.
33pub type Float32Builder = PrimitiveBuilder<Float32Type>;
34/// A 64-bit floating point array builder.
35pub type Float64Builder = PrimitiveBuilder<Float64Type>;
36
37/// A timestamp second array builder.
38pub type TimestampSecondBuilder = PrimitiveBuilder<TimestampSecondType>;
39/// A timestamp millisecond array builder.
40pub type TimestampMillisecondBuilder = PrimitiveBuilder<TimestampMillisecondType>;
41/// A timestamp microsecond array builder.
42pub type TimestampMicrosecondBuilder = PrimitiveBuilder<TimestampMicrosecondType>;
43/// A timestamp nanosecond array builder.
44pub type TimestampNanosecondBuilder = PrimitiveBuilder<TimestampNanosecondType>;
45
46/// A 32-bit date array builder.
47pub type Date32Builder = PrimitiveBuilder<Date32Type>;
48/// A 64-bit date array builder.
49pub type Date64Builder = PrimitiveBuilder<Date64Type>;
50
51/// A 32-bit elapsed time in seconds array builder.
52pub type Time32SecondBuilder = PrimitiveBuilder<Time32SecondType>;
53/// A 32-bit elapsed time in milliseconds array builder.
54pub type Time32MillisecondBuilder = PrimitiveBuilder<Time32MillisecondType>;
55/// A 64-bit elapsed time in microseconds array builder.
56pub type Time64MicrosecondBuilder = PrimitiveBuilder<Time64MicrosecondType>;
57/// A 64-bit elapsed time in nanoseconds array builder.
58pub type Time64NanosecondBuilder = PrimitiveBuilder<Time64NanosecondType>;
59
60/// A “calendar” interval in months array builder.
61pub type IntervalYearMonthBuilder = PrimitiveBuilder<IntervalYearMonthType>;
62/// A “calendar” interval in days and milliseconds array builder.
63pub type IntervalDayTimeBuilder = PrimitiveBuilder<IntervalDayTimeType>;
64/// A “calendar” interval in months, days, and nanoseconds array builder.
65pub type IntervalMonthDayNanoBuilder = PrimitiveBuilder<IntervalMonthDayNanoType>;
66
67/// An elapsed time in seconds array builder.
68pub type DurationSecondBuilder = PrimitiveBuilder<DurationSecondType>;
69/// An elapsed time in milliseconds array builder.
70pub type DurationMillisecondBuilder = PrimitiveBuilder<DurationMillisecondType>;
71/// An elapsed time in microseconds array builder.
72pub type DurationMicrosecondBuilder = PrimitiveBuilder<DurationMicrosecondType>;
73/// An elapsed time in nanoseconds array builder.
74pub type DurationNanosecondBuilder = PrimitiveBuilder<DurationNanosecondType>;
75
76/// A decimal 32 array builder
77pub type Decimal32Builder = PrimitiveBuilder<Decimal32Type>;
78/// A decimal 64 array builder
79pub type Decimal64Builder = PrimitiveBuilder<Decimal64Type>;
80/// A decimal 128 array builder
81pub type Decimal128Builder = PrimitiveBuilder<Decimal128Type>;
82/// A decimal 256 array builder
83pub type Decimal256Builder = PrimitiveBuilder<Decimal256Type>;
84
85/// Builder for [`PrimitiveArray`]
86///
87/// # Performance
88///
89/// Rust's `Vec` is highly optimized, and Arrow's conversion from `Vec` to
90/// [`PrimitiveArray`] is zero-copy — the array reuses the same underlying allocation
91/// without any data being copied. If your values are already in a `Vec`, or can be
92/// collected into one, prefer constructing a [`PrimitiveArray`] directly:
93///
94/// ```
95/// # use arrow_array::{Int32Array, Array};
96/// // Zero-copy: the array reuses the Vec's allocation
97/// let array = Int32Array::from(vec![1, 2, 3]);
98/// assert_eq!(array.len(), 3);
99/// ```
100///
101/// Internally, [`PrimitiveBuilder`] is itself backed by a `Vec<T::Native>` and a
102/// [`NullBufferBuilder`], so using one does not unlock any additional performance —
103/// it is simply a convenience wrapper for incremental construction.
104///
105/// # When to use [`PrimitiveBuilder`]
106///
107/// Prefer the builder when your array **may contain nulls but you don't know their
108/// positions upfront**. Managing a `Vec<T>` and a [`NullBufferBuilder`] in parallel
109/// by hand is error-prone; the builder keeps them in sync automatically as you call
110/// [`append_value`](PrimitiveBuilder::append_value) and
111/// [`append_null`](PrimitiveBuilder::append_null).
112///
113/// ```
114/// # use arrow_array::builder::Int32Builder;
115/// # use arrow_array::Array;
116/// let mut builder = Int32Builder::new();
117/// for (i, v) in [1, 2, 3].iter().enumerate() {
118///     if i == 1 {
119///         builder.append_null();
120///     } else {
121///         builder.append_value(*v);
122///     }
123/// }
124/// let array = builder.finish();
125/// assert_eq!(array.len(), 3);
126/// assert!(array.is_null(1));
127/// ```
128#[derive(Debug)]
129pub struct PrimitiveBuilder<T: ArrowPrimitiveType> {
130    values_builder: Vec<T::Native>,
131    null_buffer_builder: NullBufferBuilder,
132    data_type: DataType,
133}
134
135impl<T: ArrowPrimitiveType> ArrayBuilder for PrimitiveBuilder<T> {
136    /// Returns the builder as a non-mutable `Any` reference.
137    fn as_any(&self) -> &dyn Any {
138        self
139    }
140
141    /// Returns the builder as a mutable `Any` reference.
142    fn as_any_mut(&mut self) -> &mut dyn Any {
143        self
144    }
145
146    /// Returns the boxed builder as a box of `Any`.
147    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
148        self
149    }
150
151    /// Returns the number of array slots in the builder
152    fn len(&self) -> usize {
153        self.values_builder.len()
154    }
155
156    /// Builds the array and reset this builder.
157    fn finish(&mut self) -> ArrayRef {
158        Arc::new(self.finish())
159    }
160
161    /// Builds the array without resetting the builder.
162    fn finish_cloned(&self) -> ArrayRef {
163        Arc::new(self.finish_cloned())
164    }
165}
166
167impl<T: ArrowPrimitiveType> Default for PrimitiveBuilder<T> {
168    fn default() -> Self {
169        Self::new()
170    }
171}
172
173impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
174    /// Creates a new primitive array builder
175    pub fn new() -> Self {
176        Self::with_capacity(1024)
177    }
178
179    /// Creates a new primitive array builder with capacity no of items
180    pub fn with_capacity(capacity: usize) -> Self {
181        Self {
182            values_builder: Vec::with_capacity(capacity),
183            null_buffer_builder: NullBufferBuilder::new(capacity),
184            data_type: T::DATA_TYPE,
185        }
186    }
187
188    /// Creates a new primitive array builder from buffers
189    pub fn new_from_buffer(
190        values_buffer: MutableBuffer,
191        null_buffer: Option<MutableBuffer>,
192    ) -> Self {
193        let values_builder: Vec<T::Native> = ScalarBuffer::<T::Native>::from(values_buffer).into();
194
195        let null_buffer_builder = null_buffer
196            .map(|buffer| NullBufferBuilder::new_from_buffer(buffer, values_builder.len()))
197            .unwrap_or_else(|| NullBufferBuilder::new_with_len(values_builder.len()));
198
199        Self {
200            values_builder,
201            null_buffer_builder,
202            data_type: T::DATA_TYPE,
203        }
204    }
205
206    /// By default [`PrimitiveBuilder`] uses [`ArrowPrimitiveType::DATA_TYPE`] as the
207    /// data type of the generated array.
208    ///
209    /// This method allows overriding the data type, to allow specifying timezones
210    /// for [`DataType::Timestamp`] or precision and scale for [`DataType::Decimal32`],
211    /// [`DataType::Decimal64`], [`DataType::Decimal128`] and [`DataType::Decimal256`]
212    ///
213    /// # Panics
214    ///
215    /// This method panics if `data_type` is not [PrimitiveArray::is_compatible]
216    pub fn with_data_type(self, data_type: DataType) -> Self {
217        assert!(
218            PrimitiveArray::<T>::is_compatible(&data_type),
219            "incompatible data type for builder, expected {} got {}",
220            T::DATA_TYPE,
221            data_type
222        );
223        Self { data_type, ..self }
224    }
225
226    /// Returns the capacity of this builder measured in slots of type `T`
227    pub fn capacity(&self) -> usize {
228        self.values_builder.capacity()
229    }
230
231    /// Appends a value of type `T` into the builder
232    #[inline]
233    pub fn append_value(&mut self, v: T::Native) {
234        self.null_buffer_builder.append_non_null();
235        self.values_builder.push(v);
236    }
237
238    /// Appends a value of type `T` into the builder `n` times
239    #[inline]
240    pub fn append_value_n(&mut self, v: T::Native, n: usize) {
241        self.null_buffer_builder.append_n_non_nulls(n);
242        self.values_builder.extend(std::iter::repeat_n(v, n));
243    }
244
245    /// Appends a null slot into the builder
246    #[inline]
247    pub fn append_null(&mut self) {
248        self.null_buffer_builder.append_null();
249        self.values_builder.push(T::Native::default());
250    }
251
252    /// Appends `n` no. of null's into the builder
253    #[inline]
254    pub fn append_nulls(&mut self, n: usize) {
255        self.null_buffer_builder.append_n_nulls(n);
256        self.values_builder
257            .extend(std::iter::repeat_n(T::Native::default(), n));
258    }
259
260    /// Appends an `Option<T>` into the builder
261    #[inline]
262    pub fn append_option(&mut self, v: Option<T::Native>) {
263        match v {
264            None => self.append_null(),
265            Some(v) => self.append_value(v),
266        }
267    }
268
269    /// Appends a slice of type `T` into the builder
270    #[inline]
271    pub fn append_slice(&mut self, v: &[T::Native]) {
272        self.null_buffer_builder.append_n_non_nulls(v.len());
273        self.values_builder.extend_from_slice(v);
274    }
275
276    /// Appends values from a slice of type `T` and a validity boolean slice
277    ///
278    /// # Panics
279    ///
280    /// Panics if `values` and `is_valid` have different lengths
281    #[inline]
282    pub fn append_values(&mut self, values: &[T::Native], is_valid: &[bool]) {
283        assert_eq!(
284            values.len(),
285            is_valid.len(),
286            "Value and validity lengths must be equal"
287        );
288        self.null_buffer_builder.append_slice(is_valid);
289        self.values_builder.extend_from_slice(values);
290    }
291
292    /// Appends values from a iter of type `Option<T>`
293    ///
294    /// # Panics
295    ///
296    /// Panics if `values` and `is_valid` have different lengths
297    #[inline]
298    pub fn extend_from_iter_option<I: IntoIterator<Item = Option<T::Native>>>(&mut self, iter: I) {
299        let iter = iter.into_iter();
300        self.values_builder.extend(iter.map(|v| match v {
301            Some(v) => {
302                self.null_buffer_builder.append_non_null();
303                v
304            }
305            None => {
306                self.null_buffer_builder.append_null();
307                T::Native::default()
308            }
309        }));
310    }
311
312    /// Appends array values and null to this builder as is
313    /// (this means that underlying null values are copied as is).
314    ///
315    /// # Panics
316    ///
317    /// Panics if `array` and `self` data types are different
318    #[inline]
319    pub fn append_array(&mut self, array: &PrimitiveArray<T>) {
320        assert_eq!(
321            &self.data_type,
322            array.data_type(),
323            "array data type mismatch"
324        );
325
326        self.values_builder.extend_from_slice(array.values());
327        if let Some(null_buffer) = array.nulls() {
328            self.null_buffer_builder.append_buffer(null_buffer);
329        } else {
330            self.null_buffer_builder.append_n_non_nulls(array.len());
331        }
332    }
333
334    /// Appends values from a trusted length iterator.
335    ///
336    /// # Safety
337    /// This requires the iterator be a trusted length. This could instead require
338    /// the iterator implement `TrustedLen` once that is stabilized.
339    #[inline]
340    pub unsafe fn append_trusted_len_iter(&mut self, iter: impl IntoIterator<Item = T::Native>) {
341        let iter = iter.into_iter();
342        let len = iter
343            .size_hint()
344            .1
345            .expect("append_trusted_len_iter requires an upper bound");
346
347        self.null_buffer_builder.append_n_non_nulls(len);
348        self.values_builder.extend(iter);
349    }
350
351    /// Builds the [`PrimitiveArray`] and reset this builder.
352    pub fn finish(&mut self) -> PrimitiveArray<T> {
353        let len = self.len();
354        let nulls = self.null_buffer_builder.finish();
355        let builder = ArrayData::builder(self.data_type.clone())
356            .len(len)
357            .add_buffer(std::mem::take(&mut self.values_builder).into())
358            .nulls(nulls);
359
360        // SAFETY: builder is constructed from valid primitive value buffer and null buffer with matching lengths
361        let array_data = unsafe { builder.build_unchecked() };
362        PrimitiveArray::<T>::from(array_data)
363    }
364
365    /// Builds the [`PrimitiveArray`] without resetting the builder.
366    pub fn finish_cloned(&self) -> PrimitiveArray<T> {
367        let len = self.len();
368        let nulls = self.null_buffer_builder.finish_cloned();
369        let values_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
370        let builder = ArrayData::builder(self.data_type.clone())
371            .len(len)
372            .add_buffer(values_buffer)
373            .nulls(nulls);
374
375        // SAFETY: builder is constructed from valid primitive value buffer and null buffer with matching lengths
376        let array_data = unsafe { builder.build_unchecked() };
377        PrimitiveArray::<T>::from(array_data)
378    }
379
380    /// Returns the current values buffer as a slice
381    pub fn values_slice(&self) -> &[T::Native] {
382        self.values_builder.as_slice()
383    }
384
385    /// Returns the current values buffer as a mutable slice
386    pub fn values_slice_mut(&mut self) -> &mut [T::Native] {
387        self.values_builder.as_mut_slice()
388    }
389
390    /// Returns the current null buffer as a slice
391    pub fn validity_slice(&self) -> Option<&[u8]> {
392        self.null_buffer_builder.as_slice()
393    }
394
395    /// Returns the current null buffer allocated capacity, in bytes.
396    pub fn validity_capacity(&self) -> usize {
397        self.null_buffer_builder.allocated_size()
398    }
399
400    /// Returns the current null buffer as a mutable slice
401    pub fn validity_slice_mut(&mut self) -> Option<&mut [u8]> {
402        self.null_buffer_builder.as_slice_mut()
403    }
404
405    /// Returns the current values buffer and null buffer as a slice
406    pub fn slices_mut(&mut self) -> (&mut [T::Native], Option<&mut [u8]>) {
407        (
408            self.values_builder.as_mut_slice(),
409            self.null_buffer_builder.as_slice_mut(),
410        )
411    }
412}
413
414impl<P: DecimalType> PrimitiveBuilder<P> {
415    /// Sets the precision and scale
416    pub fn with_precision_and_scale(self, precision: u8, scale: i8) -> Result<Self, ArrowError> {
417        validate_decimal_precision_and_scale::<P>(precision, scale)?;
418        Ok(Self {
419            data_type: P::TYPE_CONSTRUCTOR(precision, scale),
420            ..self
421        })
422    }
423}
424
425impl<P: ArrowTimestampType> PrimitiveBuilder<P> {
426    /// Sets the timezone
427    pub fn with_timezone(self, timezone: impl Into<Arc<str>>) -> Self {
428        self.with_timezone_opt(Some(timezone.into()))
429    }
430
431    /// Sets an optional timezone
432    pub fn with_timezone_opt<S: Into<Arc<str>>>(self, timezone: Option<S>) -> Self {
433        Self {
434            data_type: DataType::Timestamp(P::UNIT, timezone.map(Into::into)),
435            ..self
436        }
437    }
438}
439
440impl<P: ArrowPrimitiveType> Extend<Option<P::Native>> for PrimitiveBuilder<P> {
441    #[inline]
442    fn extend<T: IntoIterator<Item = Option<P::Native>>>(&mut self, iter: T) {
443        for v in iter {
444            self.append_option(v)
445        }
446    }
447}
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452    use arrow_buffer::{NullBuffer, ScalarBuffer};
453    use arrow_schema::TimeUnit;
454
455    use crate::array::Array;
456    use crate::array::BooleanArray;
457    use crate::array::Date32Array;
458    use crate::array::Int32Array;
459    use crate::array::TimestampSecondArray;
460
461    #[test]
462    fn test_primitive_array_builder_i32() {
463        let mut builder = Int32Array::builder(5);
464        for i in 0..5 {
465            builder.append_value(i);
466        }
467        let arr = builder.finish();
468        assert_eq!(5, arr.len());
469        assert_eq!(0, arr.offset());
470        assert_eq!(0, arr.null_count());
471        for i in 0..5 {
472            assert!(!arr.is_null(i));
473            assert!(arr.is_valid(i));
474            assert_eq!(i as i32, arr.value(i));
475        }
476    }
477
478    #[test]
479    fn test_primitive_array_builder_i32_append_iter() {
480        let mut builder = Int32Array::builder(5);
481        unsafe { builder.append_trusted_len_iter(0..5) };
482        let arr = builder.finish();
483        assert_eq!(5, arr.len());
484        assert_eq!(0, arr.offset());
485        assert_eq!(0, arr.null_count());
486        for i in 0..5 {
487            assert!(!arr.is_null(i));
488            assert!(arr.is_valid(i));
489            assert_eq!(i as i32, arr.value(i));
490        }
491    }
492
493    #[test]
494    fn test_primitive_array_builder_i32_append_nulls() {
495        let mut builder = Int32Array::builder(5);
496        builder.append_nulls(5);
497        let arr = builder.finish();
498        assert_eq!(5, arr.len());
499        assert_eq!(0, arr.offset());
500        assert_eq!(5, arr.null_count());
501        for i in 0..5 {
502            assert!(arr.is_null(i));
503            assert!(!arr.is_valid(i));
504        }
505    }
506
507    #[test]
508    fn test_primitive_array_builder_date32() {
509        let mut builder = Date32Array::builder(5);
510        for i in 0..5 {
511            builder.append_value(i);
512        }
513        let arr = builder.finish();
514        assert_eq!(5, arr.len());
515        assert_eq!(0, arr.offset());
516        assert_eq!(0, arr.null_count());
517        for i in 0..5 {
518            assert!(!arr.is_null(i));
519            assert!(arr.is_valid(i));
520            assert_eq!(i as i32, arr.value(i));
521        }
522    }
523
524    #[test]
525    fn test_primitive_array_builder_timestamp_second() {
526        let mut builder = TimestampSecondArray::builder(5);
527        for i in 0..5 {
528            builder.append_value(i);
529        }
530        let arr = builder.finish();
531        assert_eq!(5, arr.len());
532        assert_eq!(0, arr.offset());
533        assert_eq!(0, arr.null_count());
534        for i in 0..5 {
535            assert!(!arr.is_null(i));
536            assert!(arr.is_valid(i));
537            assert_eq!(i as i64, arr.value(i));
538        }
539    }
540
541    #[test]
542    fn test_primitive_array_builder_bool() {
543        // 00000010 01001000
544        let buf = Buffer::from([72_u8, 2_u8]);
545        let mut builder = BooleanArray::builder(10);
546        for i in 0..10 {
547            if i == 3 || i == 6 || i == 9 {
548                builder.append_value(true);
549            } else {
550                builder.append_value(false);
551            }
552        }
553
554        let arr = builder.finish();
555        assert_eq!(&buf, arr.values().inner());
556        assert_eq!(10, arr.len());
557        assert_eq!(0, arr.offset());
558        assert_eq!(0, arr.null_count());
559        for i in 0..10 {
560            assert!(!arr.is_null(i));
561            assert!(arr.is_valid(i));
562            assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {i}")
563        }
564    }
565
566    #[test]
567    fn test_primitive_array_builder_append_option() {
568        let arr1 = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
569
570        let mut builder = Int32Array::builder(5);
571        builder.append_option(Some(0));
572        builder.append_option(None);
573        builder.append_option(Some(2));
574        builder.append_option(None);
575        builder.append_option(Some(4));
576        let arr2 = builder.finish();
577
578        assert_eq!(arr1.len(), arr2.len());
579        assert_eq!(arr1.offset(), arr2.offset());
580        assert_eq!(arr1.null_count(), arr2.null_count());
581        for i in 0..5 {
582            assert_eq!(arr1.is_null(i), arr2.is_null(i));
583            assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
584            if arr1.is_valid(i) {
585                assert_eq!(arr1.value(i), arr2.value(i));
586            }
587        }
588    }
589
590    #[test]
591    fn test_primitive_array_builder_append_null() {
592        let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]);
593
594        let mut builder = Int32Array::builder(5);
595        builder.append_value(0);
596        builder.append_value(2);
597        builder.append_null();
598        builder.append_null();
599        builder.append_value(4);
600        let arr2 = builder.finish();
601
602        assert_eq!(arr1.len(), arr2.len());
603        assert_eq!(arr1.offset(), arr2.offset());
604        assert_eq!(arr1.null_count(), arr2.null_count());
605        for i in 0..5 {
606            assert_eq!(arr1.is_null(i), arr2.is_null(i));
607            assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
608            if arr1.is_valid(i) {
609                assert_eq!(arr1.value(i), arr2.value(i));
610            }
611        }
612    }
613
614    #[test]
615    fn test_primitive_array_builder_append_slice() {
616        let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]);
617
618        let mut builder = Int32Array::builder(5);
619        builder.append_slice(&[0, 2]);
620        builder.append_null();
621        builder.append_null();
622        builder.append_value(4);
623        let arr2 = builder.finish();
624
625        assert_eq!(arr1.len(), arr2.len());
626        assert_eq!(arr1.offset(), arr2.offset());
627        assert_eq!(arr1.null_count(), arr2.null_count());
628        for i in 0..5 {
629            assert_eq!(arr1.is_null(i), arr2.is_null(i));
630            assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
631            if arr1.is_valid(i) {
632                assert_eq!(arr1.value(i), arr2.value(i));
633            }
634        }
635    }
636
637    #[test]
638    fn test_primitive_array_builder_finish() {
639        let mut builder = Int32Builder::new();
640        builder.append_slice(&[2, 4, 6, 8]);
641        let mut arr = builder.finish();
642        assert_eq!(4, arr.len());
643        assert_eq!(0, builder.len());
644
645        builder.append_slice(&[1, 3, 5, 7, 9]);
646        arr = builder.finish();
647        assert_eq!(5, arr.len());
648        assert_eq!(0, builder.len());
649    }
650
651    #[test]
652    fn test_primitive_array_builder_finish_cloned() {
653        let mut builder = Int32Builder::new();
654        builder.append_value(23);
655        builder.append_value(45);
656        let result = builder.finish_cloned();
657        assert_eq!(result, Int32Array::from(vec![23, 45]));
658        builder.append_value(56);
659        assert_eq!(builder.finish_cloned(), Int32Array::from(vec![23, 45, 56]));
660
661        builder.append_slice(&[2, 4, 6, 8]);
662        let mut arr = builder.finish();
663        assert_eq!(7, arr.len());
664        assert_eq!(arr, Int32Array::from(vec![23, 45, 56, 2, 4, 6, 8]));
665        assert_eq!(0, builder.len());
666
667        builder.append_slice(&[1, 3, 5, 7, 9]);
668        arr = builder.finish();
669        assert_eq!(5, arr.len());
670        assert_eq!(0, builder.len());
671    }
672
673    #[test]
674    fn test_primitive_array_builder_with_data_type() {
675        let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2));
676        builder.append_value(1);
677        let array = builder.finish();
678        assert_eq!(array.precision(), 1);
679        assert_eq!(array.scale(), 2);
680
681        let data_type = DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into()));
682        let mut builder = TimestampNanosecondBuilder::new().with_data_type(data_type.clone());
683        builder.append_value(1);
684        let array = builder.finish();
685        assert_eq!(array.data_type(), &data_type);
686    }
687
688    #[test]
689    #[should_panic(expected = "incompatible data type for builder, expected Int32 got Int64")]
690    fn test_invalid_with_data_type() {
691        Int32Builder::new().with_data_type(DataType::Int64);
692    }
693
694    #[test]
695    fn test_extend() {
696        let mut builder = PrimitiveBuilder::<Int16Type>::new();
697        builder.extend([1, 2, 3, 5, 2, 4, 4].into_iter().map(Some));
698        builder.extend([2, 4, 6, 2].into_iter().map(Some));
699        let array = builder.finish();
700        assert_eq!(array.values(), &[1, 2, 3, 5, 2, 4, 4, 2, 4, 6, 2]);
701    }
702
703    #[test]
704    fn test_primitive_array_append_array() {
705        let input = vec![
706            Some(1),
707            None,
708            Some(3),
709            None,
710            Some(5),
711            None,
712            None,
713            None,
714            Some(7),
715            Some(9),
716            Some(8),
717            Some(6),
718            Some(4),
719        ];
720        let arr1 = Int32Array::from(input[..5].to_vec());
721        let arr2 = Int32Array::from(input[5..8].to_vec());
722        let arr3 = Int32Array::from(input[8..].to_vec());
723
724        let mut builder = Int32Array::builder(5);
725        builder.append_array(&arr1);
726        builder.append_array(&arr2);
727        builder.append_array(&arr3);
728        let actual = builder.finish();
729        let expected = Int32Array::from(input);
730
731        assert_eq!(actual, expected);
732    }
733
734    #[test]
735    fn test_append_array_add_underlying_null_values() {
736        let array = Int32Array::new(
737            ScalarBuffer::from(vec![2, 3, 4, 5]),
738            Some(NullBuffer::from(&[true, true, false, false])),
739        );
740
741        let mut builder = Int32Array::builder(5);
742        builder.append_array(&array);
743        let actual = builder.finish();
744
745        assert_eq!(actual, array);
746        assert_eq!(actual.values(), array.values())
747    }
748
749    #[test]
750    #[should_panic(expected = "array data type mismatch")]
751    fn test_invalid_with_data_type_in_append_array() {
752        let array = {
753            let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2));
754            builder.append_value(1);
755            builder.finish()
756        };
757
758        let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(2, 3));
759        builder.append_array(&array)
760    }
761}
arrow_array/builder/primitive_builder.rs

arrow_array/builder/
primitive_builder.rs
