use std::sync::Arc;
use arrow_array::ArrayRef as ArrowArrayRef;
use arrow_array::Decimal32Array as ArrowDecimal32Array;
use arrow_array::Decimal64Array as ArrowDecimal64Array;
use arrow_array::Decimal128Array as ArrowDecimal128Array;
use arrow_array::Decimal256Array as ArrowDecimal256Array;
use arrow_buffer::i256;
use arrow_schema::DataType;
use itertools::Itertools;
use num_traits::AsPrimitive;
use num_traits::ToPrimitive;
use vortex_array::ArrayRef;
use vortex_array::ExecutionCtx;
use vortex_array::arrays::DecimalArray;
use vortex_array::dtype::DType;
use vortex_array::dtype::DecimalType;
use vortex_buffer::Buffer;
use vortex_error::VortexResult;
use vortex_error::vortex_ensure;
use vortex_error::vortex_err;
use crate::null_buffer::to_null_buffer;
pub(super) fn to_arrow_decimal(
array: ArrayRef,
data_type: &DataType,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrowArrayRef> {
vortex_ensure!(
matches!(array.dtype(), DType::Decimal(..)),
"Cannot convert Vortex array with dtype {} to an Arrow {} array",
array.dtype(),
data_type
);
let decimal_array = array.execute::<DecimalArray>(ctx)?;
match data_type {
DataType::Decimal32(..) => to_arrow_decimal32(decimal_array, ctx),
DataType::Decimal64(..) => to_arrow_decimal64(decimal_array, ctx),
DataType::Decimal128(..) => to_arrow_decimal128(decimal_array, ctx),
DataType::Decimal256(..) => to_arrow_decimal256(decimal_array, ctx),
_ => unreachable!("to_arrow_decimal called with non-decimal type"),
}
}
fn to_arrow_decimal32(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult<ArrowArrayRef> {
let null_buffer = to_null_buffer(
array
.as_ref()
.validity()?
.execute_mask(array.as_ref().len(), ctx)?,
);
let buffer: Buffer<i32> = match array.values_type() {
DecimalType::I8 => {
Buffer::from_trusted_len_iter(array.buffer::<i8>().into_iter().map(|x| x.as_()))
}
DecimalType::I16 => {
Buffer::from_trusted_len_iter(array.buffer::<i16>().into_iter().map(|x| x.as_()))
}
DecimalType::I32 => array.buffer::<i32>(),
DecimalType::I64 => array
.buffer::<i64>()
.into_iter()
.map(|x| {
x.to_i32()
.ok_or_else(|| vortex_err!("i64 to i32 narrowing cannot be done safely"))
})
.process_results(|iter| Buffer::from_trusted_len_iter(iter))?,
DecimalType::I128 => array
.buffer::<i128>()
.into_iter()
.map(|x| {
x.to_i32()
.ok_or_else(|| vortex_err!("i128 to i32 narrowing cannot be done safely"))
})
.process_results(|iter| Buffer::from_trusted_len_iter(iter))?,
DecimalType::I256 => array
.buffer::<vortex_array::dtype::i256>()
.into_iter()
.map(|x| {
x.to_i32()
.ok_or_else(|| vortex_err!("i256 to i32 narrowing cannot be done safely"))
})
.process_results(|iter| Buffer::from_trusted_len_iter(iter))?,
};
Ok(Arc::new(
ArrowDecimal32Array::new(buffer.into_arrow_scalar_buffer(), null_buffer)
.with_precision_and_scale(
array.decimal_dtype().precision(),
array.decimal_dtype().scale(),
)?,
))
}
fn to_arrow_decimal64(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult<ArrowArrayRef> {
let null_buffer = to_null_buffer(
array
.as_ref()
.validity()?
.execute_mask(array.as_ref().len(), ctx)?,
);
let buffer: Buffer<i64> = match array.values_type() {
DecimalType::I8 => {
Buffer::from_trusted_len_iter(array.buffer::<i8>().into_iter().map(|x| x.as_()))
}
DecimalType::I16 => {
Buffer::from_trusted_len_iter(array.buffer::<i16>().into_iter().map(|x| x.as_()))
}
DecimalType::I32 => {
Buffer::from_trusted_len_iter(array.buffer::<i32>().into_iter().map(|x| x.as_()))
}
DecimalType::I64 => array.buffer::<i64>(),
DecimalType::I128 => array
.buffer::<i128>()
.into_iter()
.map(|x| {
x.to_i64()
.ok_or_else(|| vortex_err!("i128 to i64 narrowing cannot be done safely"))
})
.process_results(|iter| Buffer::from_trusted_len_iter(iter))?,
DecimalType::I256 => array
.buffer::<vortex_array::dtype::i256>()
.into_iter()
.map(|x| {
x.to_i64()
.ok_or_else(|| vortex_err!("i256 to i64 narrowing cannot be done safely"))
})
.process_results(|iter| Buffer::from_trusted_len_iter(iter))?,
};
Ok(Arc::new(
ArrowDecimal64Array::new(buffer.into_arrow_scalar_buffer(), null_buffer)
.with_precision_and_scale(
array.decimal_dtype().precision(),
array.decimal_dtype().scale(),
)?,
))
}
fn to_arrow_decimal128(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult<ArrowArrayRef> {
let null_buffer = to_null_buffer(
array
.as_ref()
.validity()?
.execute_mask(array.as_ref().len(), ctx)?,
);
let buffer: Buffer<i128> = match array.values_type() {
DecimalType::I8 => {
Buffer::from_trusted_len_iter(array.buffer::<i8>().into_iter().map(|x| x.as_()))
}
DecimalType::I16 => {
Buffer::from_trusted_len_iter(array.buffer::<i16>().into_iter().map(|x| x.as_()))
}
DecimalType::I32 => {
Buffer::from_trusted_len_iter(array.buffer::<i32>().into_iter().map(|x| x.as_()))
}
DecimalType::I64 => {
Buffer::from_trusted_len_iter(array.buffer::<i64>().into_iter().map(|x| x.as_()))
}
DecimalType::I128 => array.buffer::<i128>(),
DecimalType::I256 => array
.buffer::<vortex_array::dtype::i256>()
.into_iter()
.map(|x| {
x.to_i128()
.ok_or_else(|| vortex_err!("i256 to i128 narrowing cannot be done safely"))
})
.process_results(|iter| Buffer::from_trusted_len_iter(iter))?,
};
Ok(Arc::new(
ArrowDecimal128Array::new(buffer.into_arrow_scalar_buffer(), null_buffer)
.with_precision_and_scale(
array.decimal_dtype().precision(),
array.decimal_dtype().scale(),
)?,
))
}
fn to_arrow_decimal256(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult<ArrowArrayRef> {
let null_buffer = to_null_buffer(
array
.as_ref()
.validity()?
.execute_mask(array.as_ref().len(), ctx)?,
);
let buffer: Buffer<i256> = match array.values_type() {
DecimalType::I8 => {
Buffer::from_trusted_len_iter(array.buffer::<i8>().into_iter().map(|x| x.as_()))
}
DecimalType::I16 => {
Buffer::from_trusted_len_iter(array.buffer::<i16>().into_iter().map(|x| x.as_()))
}
DecimalType::I32 => {
Buffer::from_trusted_len_iter(array.buffer::<i32>().into_iter().map(|x| x.as_()))
}
DecimalType::I64 => {
Buffer::from_trusted_len_iter(array.buffer::<i64>().into_iter().map(|x| x.as_()))
}
DecimalType::I128 => Buffer::from_trusted_len_iter(
array
.buffer::<i128>()
.into_iter()
.map(|x| vortex_array::dtype::i256::from_i128(x).into()),
),
DecimalType::I256 => {
Buffer::<i256>::from_byte_buffer(array.buffer_handle().clone().into_host_sync())
}
};
Ok(Arc::new(
ArrowDecimal256Array::new(buffer.into_arrow_scalar_buffer(), null_buffer)
.with_precision_and_scale(
array.decimal_dtype().precision(),
array.decimal_dtype().scale(),
)?,
))
}
#[cfg(test)]
mod tests {
use arrow_array::Array;
use arrow_array::Decimal128Array;
use arrow_array::Decimal256Array;
use arrow_buffer::i256;
use arrow_schema::DataType;
use rstest::rstest;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::array_session;
use vortex_array::builders::ArrayBuilder;
use vortex_array::builders::DecimalBuilder;
use vortex_array::dtype::DecimalDType;
use vortex_array::dtype::NativeDecimalType;
use vortex_array::validity::Validity;
use vortex_buffer::buffer;
use vortex_error::VortexResult;
use crate::ArrowArrayExecutor;
use crate::executor::decimal::DecimalArray;
#[test]
fn decimal_to_arrow() -> VortexResult<()> {
let mut ctx = array_session().create_execution_ctx();
let decimal_vortex = DecimalArray::new(
buffer![1i128, 2i128, 3i128, 4i128, 5i128],
DecimalDType::new(19, 2),
Validity::NonNullable,
);
let arrow = decimal_vortex
.into_array()
.execute_arrow(Some(&DataType::Decimal128(19, 2)), &mut ctx)?;
assert_eq!(arrow.data_type(), &DataType::Decimal128(19, 2));
let decimal_array = arrow.as_any().downcast_ref::<Decimal128Array>().unwrap();
assert_eq!(
decimal_array.values().as_ref(),
&[1i128, 2i128, 3i128, 4i128, 5i128]
);
assert_eq!(
decimal_array,
&Decimal128Array::from(vec![1, 2, 3, 4, 5]).with_precision_and_scale(19, 2)?
);
Ok(())
}
#[rstest]
#[case(0i8)]
#[case(0i16)]
#[case(0i32)]
#[case(0i64)]
#[case(0i128)]
#[case(vortex_array::dtype::i256::ZERO)]
fn test_to_arrow_decimal128<T: NativeDecimalType>(
#[case] _decimal_type: T,
) -> VortexResult<()> {
let mut ctx = array_session().create_execution_ctx();
let mut decimal = DecimalBuilder::new_in::<T>(
DecimalDType::new(2, 1),
false.into(),
vortex_buffer::BufferAllocatorRef::static_ref(),
);
decimal.append_value(10);
decimal.append_value(11);
decimal.append_value(12);
let decimal = decimal.finish();
let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal128(2, 1)), &mut ctx)?;
let arrow_decimal = arrow_array
.as_any()
.downcast_ref::<Decimal128Array>()
.unwrap();
assert_eq!(arrow_decimal.value(0), 10);
assert_eq!(arrow_decimal.value(1), 11);
assert_eq!(arrow_decimal.value(2), 12);
assert_eq!(
arrow_decimal,
&Decimal128Array::from(vec![10, 11, 12]).with_precision_and_scale(2, 1)?
);
Ok(())
}
#[rstest]
#[case(0i8)]
#[case(0i16)]
#[case(0i32)]
#[case(0i64)]
#[case(0i128)]
#[case(vortex_array::dtype::i256::ZERO)]
fn test_to_arrow_decimal32<T: NativeDecimalType>(#[case] _decimal_type: T) -> VortexResult<()> {
use arrow_array::Decimal32Array;
let mut ctx = array_session().create_execution_ctx();
let mut decimal = DecimalBuilder::new_in::<T>(
DecimalDType::new(2, 1),
false.into(),
vortex_buffer::BufferAllocatorRef::static_ref(),
);
decimal.append_value(10);
decimal.append_value(11);
decimal.append_value(12);
let decimal = decimal.finish();
let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal32(2, 1)), &mut ctx)?;
let arrow_decimal = arrow_array
.as_any()
.downcast_ref::<Decimal32Array>()
.unwrap();
assert_eq!(arrow_decimal.value(0), 10);
assert_eq!(arrow_decimal.value(1), 11);
assert_eq!(arrow_decimal.value(2), 12);
assert_eq!(
arrow_decimal,
&Decimal32Array::from(vec![10, 11, 12]).with_precision_and_scale(2, 1)?
);
Ok(())
}
#[rstest]
#[case(0i8)]
#[case(0i16)]
#[case(0i32)]
#[case(0i64)]
#[case(0i128)]
#[case(vortex_array::dtype::i256::ZERO)]
fn test_to_arrow_decimal64<T: NativeDecimalType>(#[case] _decimal_type: T) -> VortexResult<()> {
use arrow_array::Decimal64Array;
let mut ctx = array_session().create_execution_ctx();
let mut decimal = DecimalBuilder::new_in::<T>(
DecimalDType::new(2, 1),
false.into(),
vortex_buffer::BufferAllocatorRef::static_ref(),
);
decimal.append_value(10);
decimal.append_value(11);
decimal.append_value(12);
let decimal = decimal.finish();
let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal64(2, 1)), &mut ctx)?;
let arrow_decimal = arrow_array
.as_any()
.downcast_ref::<Decimal64Array>()
.unwrap();
assert_eq!(arrow_decimal.value(0), 10);
assert_eq!(arrow_decimal.value(1), 11);
assert_eq!(arrow_decimal.value(2), 12);
assert_eq!(
arrow_decimal,
&Decimal64Array::from(vec![10, 11, 12]).with_precision_and_scale(2, 1)?
);
Ok(())
}
#[rstest]
#[case(0i8)]
#[case(0i16)]
#[case(0i32)]
#[case(0i64)]
#[case(0i128)]
#[case(vortex_array::dtype::i256::ZERO)]
fn test_to_arrow_decimal256<T: NativeDecimalType>(
#[case] _decimal_type: T,
) -> VortexResult<()> {
let mut ctx = array_session().create_execution_ctx();
let mut decimal = DecimalBuilder::new_in::<T>(
DecimalDType::new(2, 1),
false.into(),
vortex_buffer::BufferAllocatorRef::static_ref(),
);
decimal.append_value(10);
decimal.append_value(11);
decimal.append_value(12);
let decimal = decimal.finish();
let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal256(2, 1)), &mut ctx)?;
let arrow_decimal = arrow_array
.as_any()
.downcast_ref::<Decimal256Array>()
.unwrap();
assert_eq!(arrow_decimal.value(0), i256::from_i128(10));
assert_eq!(arrow_decimal.value(1), i256::from_i128(11));
assert_eq!(arrow_decimal.value(2), i256::from_i128(12));
assert_eq!(
arrow_decimal,
&Decimal256Array::from(vec![
i256::from_i128(10),
i256::from_i128(11),
i256::from_i128(12)
])
.with_precision_and_scale(2, 1)?
);
Ok(())
}
}