diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index 3199104382a6..febafcc6f02a 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -269,24 +269,55 @@ impl PrimitiveArray { /// /// # Panics /// - /// Panics if: - /// - `values.len() != nulls.len()` - /// - `!Self::is_compatible(data_type)` + /// Panics if [`Self::try_new`] returns an error pub fn new( data_type: DataType, values: ScalarBuffer, nulls: Option, ) -> Self { - Self::assert_compatible(&data_type); + Self::try_new(data_type, values, nulls).unwrap() + } + + /// Create a new [`PrimitiveArray`] from the provided data_type, values, nulls + /// + /// # Errors + /// + /// Errors if: + /// - `values.len() != nulls.len()` + /// - `!Self::is_compatible(data_type)` + pub fn try_new( + data_type: DataType, + values: ScalarBuffer, + nulls: Option, + ) -> Result { + if !Self::is_compatible(&data_type) { + return Err(ArrowError::InvalidArgumentError(format!( + "PrimitiveArray expected data type {} got {}", + T::DATA_TYPE, + data_type + ))); + } + if let Some(n) = nulls.as_ref() { - assert_eq!(values.len(), n.len()); + if n.len() != values.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Incorrect length of null buffer for PrimitiveArray, expected {} got {}", + values.len(), + n.len(), + ))); + } } - Self { + Ok(Self { data_type, values, nulls, - } + }) + } + + /// Deconstruct this array into its constituent parts + pub fn into_parts(self) -> (DataType, ScalarBuffer, Option) { + (self.data_type, self.values, self.nulls) } /// Asserts that `data_type` is compatible with `Self` @@ -2262,4 +2293,33 @@ mod tests { let array = array.with_timezone("+02:00"); assert_eq!(array.timezone(), Some("+02:00")); } + + #[test] + fn test_try_new() { + Int32Array::new(DataType::Int32, vec![1, 2, 3, 4].into(), None); + Int32Array::new( + DataType::Int32, + vec![1, 2, 3, 4].into(), + Some(NullBuffer::new_null(4)), + ); + let err = Int32Array::try_new(DataType::Int64, vec![1, 2, 3, 4].into(), None) + .unwrap_err(); + + assert_eq!( + err.to_string(), + "Invalid argument error: PrimitiveArray expected data type Int32 got Int64" + ); + + let err = Int32Array::try_new( + DataType::Int32, + vec![1, 2, 3, 4].into(), + Some(NullBuffer::new_null(3)), + ) + .unwrap_err(); + + assert_eq!( + err.to_string(), + "Invalid argument error: Incorrect length of null buffer for PrimitiveArray, expected 4 got 3" + ); + } }