Avoid checking DataType when unecessary

Author: alamb

arrow-array/src/array/primitive_array.rs

@@ -638,22 +638,6 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
         Self::try_new(values, nulls).unwrap()
     }
 
-    /// Create a [`PrimitiveArray`] from the provided data type, values and nulls
-    ///
-    /// panic's if the type is not compatible
-    pub(crate) fn new_with_datatype(
-        data_type: DataType,
-        values: ScalarBuffer<T::Native>,
-        nulls: Option<NullBuffer>,
-    ) -> Self {
-        Self::assert_compatible(&data_type);
-        Self {
-            data_type,
-            values,
-            nulls,
-        }
-    }
-
     /// Create a new [`PrimitiveArray`] of the given length where all values are null
     pub fn new_null(length: usize) -> Self {
         Self {

arrow-array/src/builder/primitive_builder.rs

@@ -99,7 +99,8 @@ pub type Decimal256Builder = PrimitiveBuilder<Decimal256Type>;
 pub struct PrimitiveBuilder<T: ArrowPrimitiveType> {
     values_builder: Vec<T::Native>,
     null_buffer_builder: NullBufferBuilder,
-    data_type: DataType,
+    /// Optional data type override (e.g. to add timezone or precision/scale)
+    data_type: Option<DataType>,
 }
 
 impl<T: ArrowPrimitiveType> ArrayBuilder for PrimitiveBuilder<T> {
@@ -151,7 +152,7 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
         Self {
             values_builder: Vec::with_capacity(capacity),
             null_buffer_builder: NullBufferBuilder::new(capacity),
-            data_type: T::DATA_TYPE,
+            data_type: None,
         }
     }
 
@@ -169,7 +170,7 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
         Self {
             values_builder,
             null_buffer_builder,
-            data_type: T::DATA_TYPE,
+            data_type: None,
         }
     }
 
@@ -190,7 +191,10 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
             T::DATA_TYPE,
             data_type
         );
-        Self { data_type, ..self }
+        Self {
+            data_type: Some(data_type),
+            ..self
+        }
     }
 
     /// Returns the capacity of this builder measured in slots of type `T`
@@ -267,11 +271,11 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
     /// Panics if `array` and `self` data types are different
     #[inline]
     pub fn append_array(&mut self, array: &PrimitiveArray<T>) {
-        assert_eq!(
-            &self.data_type,
-            array.data_type(),
-            "array data type mismatch"
-        );
+        if let Some(data_type) = &self.data_type {
+            assert_eq!(data_type, array.data_type(), "array data type mismatch");
+        } else {
+            assert_eq!(&T::DATA_TYPE, array.data_type(), "array data type mismatch")
+        }
 
         self.values_builder.extend_from_slice(array.values());
         if let Some(null_buffer) = array.nulls() {
@@ -304,22 +308,22 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
     pub fn finish(&mut self) -> PrimitiveArray<T> {
         let values_buffer = Buffer::from(std::mem::take(&mut self.values_builder));
         let nulls = self.null_buffer_builder.finish();
-        PrimitiveArray::<T>::new_with_datatype(
-            self.data_type.clone(),
-            ScalarBuffer::from(values_buffer),
-            nulls,
-        )
+        let mut new_array = PrimitiveArray::<T>::new(ScalarBuffer::from(values_buffer), nulls);
+        if let Some(data_type) = &self.data_type {
+            new_array = new_array.with_data_type(data_type.clone())
+        }
+        new_array
     }
 
     /// Builds the [`PrimitiveArray`] without resetting the builder.
     pub fn finish_cloned(&self) -> PrimitiveArray<T> {
         let nulls = self.null_buffer_builder.finish_cloned();
         let values_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
-        PrimitiveArray::<T>::new_with_datatype(
-            self.data_type.clone(),
-            ScalarBuffer::from(values_buffer),
-            nulls,
-        )
+        let mut new_array = PrimitiveArray::<T>::new(ScalarBuffer::from(values_buffer), nulls);
+        if let Some(data_type) = &self.data_type {
+            new_array = new_array.with_data_type(data_type.clone())
+        }
+        new_array
     }
 
     /// Builds the [`PrimitiveArray`] , consuming this builder.
@@ -344,11 +348,14 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
             null_buffer_builder,
             data_type,
         } = self;
-        PrimitiveArray::new_with_datatype(
-            data_type, // reuse data_type to save a DataType::drop()
-            ScalarBuffer::from(values_builder),
-            null_buffer_builder.build(),
-        )
+
+        let nulls = null_buffer_builder.build();
+        let values_buffer = ScalarBuffer::from(Buffer::from(values_builder));
+        let mut new_array = PrimitiveArray::<T>::new(values_buffer, nulls);
+        if let Some(data_type) = data_type {
+            new_array = new_array.with_data_type(data_type)
+        }
+        new_array
     }
 
     /// Returns the current values buffer as a slice
@@ -385,7 +392,7 @@ impl<P: DecimalType> PrimitiveBuilder<P> {
     pub fn with_precision_and_scale(self, precision: u8, scale: i8) -> Result<Self, ArrowError> {
         validate_decimal_precision_and_scale::<P>(precision, scale)?;
         Ok(Self {
-            data_type: P::TYPE_CONSTRUCTOR(precision, scale),
+            data_type: Some(P::TYPE_CONSTRUCTOR(precision, scale)),
             ..self
         })
     }
@@ -400,7 +407,7 @@ impl<P: ArrowTimestampType> PrimitiveBuilder<P> {
     /// Sets an optional timezone
     pub fn with_timezone_opt<S: Into<Arc<str>>>(self, timezone: Option<S>) -> Self {
         Self {
-            data_type: DataType::Timestamp(P::UNIT, timezone.map(Into::into)),
+            data_type: Some(DataType::Timestamp(P::UNIT, timezone.map(Into::into))),
             ..self
         }
     }