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

arrow_array/builder/
primitive_builder.rs
