Skip to content

Commit

Permalink
Supplement as_*_array functions
Browse files Browse the repository at this point in the history
These functions are useful when implementing UDFs. It would be helpful
for downstream projects to provide graceful cast functions for all data
types.
  • Loading branch information
findepi committed Nov 27, 2024
1 parent 6a4bf0f commit 612a35f
Showing 1 changed file with 38 additions and 16 deletions.
54 changes: 38 additions & 16 deletions datafusion/common/src/cast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,33 +36,65 @@ use arrow::{
},
datatypes::{ArrowDictionaryKeyType, ArrowPrimitiveType},
};
use arrow_array::{BinaryViewArray, StringViewArray};
use arrow_array::{
BinaryViewArray, Float16Array, Int16Array, Int8Array, StringViewArray, UInt16Array,
};

// Downcast ArrayRef to Date32Array
pub fn as_date32_array(array: &dyn Array) -> Result<&Date32Array> {
Ok(downcast_value!(array, Date32Array))
}

// Downcast ArrayRef to Date64Array
pub fn as_date64_array(array: &dyn Array) -> Result<&Date64Array> {
Ok(downcast_value!(array, Date64Array))
}

// Downcast ArrayRef to StructArray
pub fn as_struct_array(array: &dyn Array) -> Result<&StructArray> {
Ok(downcast_value!(array, StructArray))
}

// Downcast ArrayRef to Int8Array
pub fn as_int8_array(array: &dyn Array) -> Result<&Int8Array> {
Ok(downcast_value!(array, Int8Array))
}

// Downcast ArrayRef to UInt8Array
pub fn as_uint8_array(array: &dyn Array) -> Result<&UInt8Array> {
Ok(downcast_value!(array, UInt8Array))
}

// Downcast ArrayRef to Int16Array
pub fn as_int16_array(array: &dyn Array) -> Result<&Int16Array> {
Ok(downcast_value!(array, Int16Array))
}

// Downcast ArrayRef to UInt16Array
pub fn as_uint16_array(array: &dyn Array) -> Result<&UInt16Array> {
Ok(downcast_value!(array, UInt16Array))
}

// Downcast ArrayRef to Int32Array
pub fn as_int32_array(array: &dyn Array) -> Result<&Int32Array> {
Ok(downcast_value!(array, Int32Array))
}

// Downcast ArrayRef to UInt32Array
pub fn as_uint32_array(array: &dyn Array) -> Result<&UInt32Array> {
Ok(downcast_value!(array, UInt32Array))
}

// Downcast ArrayRef to Int64Array
pub fn as_int64_array(array: &dyn Array) -> Result<&Int64Array> {
Ok(downcast_value!(array, Int64Array))
}

// Downcast ArrayRef to UInt64Array
pub fn as_uint64_array(array: &dyn Array) -> Result<&UInt64Array> {
Ok(downcast_value!(array, UInt64Array))
}

// Downcast ArrayRef to Decimal128Array
pub fn as_decimal128_array(array: &dyn Array) -> Result<&Decimal128Array> {
Ok(downcast_value!(array, Decimal128Array))
Expand All @@ -73,6 +105,11 @@ pub fn as_decimal256_array(array: &dyn Array) -> Result<&Decimal256Array> {
Ok(downcast_value!(array, Decimal256Array))
}

// Downcast ArrayRef to Float16Array
pub fn as_float16_array(array: &dyn Array) -> Result<&Float16Array> {
Ok(downcast_value!(array, Float16Array))
}

// Downcast ArrayRef to Float32Array
pub fn as_float32_array(array: &dyn Array) -> Result<&Float32Array> {
Ok(downcast_value!(array, Float32Array))
Expand All @@ -93,16 +130,6 @@ pub fn as_string_view_array(array: &dyn Array) -> Result<&StringViewArray> {
Ok(downcast_value!(array, StringViewArray))
}

// Downcast ArrayRef to UInt32Array
pub fn as_uint32_array(array: &dyn Array) -> Result<&UInt32Array> {
Ok(downcast_value!(array, UInt32Array))
}

// Downcast ArrayRef to UInt64Array
pub fn as_uint64_array(array: &dyn Array) -> Result<&UInt64Array> {
Ok(downcast_value!(array, UInt64Array))
}

// Downcast ArrayRef to BooleanArray
pub fn as_boolean_array(array: &dyn Array) -> Result<&BooleanArray> {
Ok(downcast_value!(array, BooleanArray))
Expand Down Expand Up @@ -242,11 +269,6 @@ pub fn as_fixed_size_binary_array(array: &dyn Array) -> Result<&FixedSizeBinaryA
Ok(downcast_value!(array, FixedSizeBinaryArray))
}

// Downcast ArrayRef to Date64Array
pub fn as_date64_array(array: &dyn Array) -> Result<&Date64Array> {
Ok(downcast_value!(array, Date64Array))
}

// Downcast ArrayRef to GenericBinaryArray
pub fn as_generic_string_array<T: OffsetSizeTrait>(
array: &dyn Array,
Expand Down

0 comments on commit 612a35f

Please sign in to comment.