use arrow2::{
buffer::Buffer,
datatypes::{DataType, Field},
types::NativeType,
};
use chrono::{NaiveDate, NaiveDateTime};
pub trait ArrowField {
type Type;
fn data_type() -> DataType;
#[inline]
#[doc(hidden)]
fn field(name: &str) -> Field {
Field::new(name, Self::data_type(), Self::is_nullable())
}
#[inline]
#[doc(hidden)]
fn is_nullable() -> bool {
false
}
}
#[macro_export]
macro_rules! arrow_enable_vec_for_type {
($t:ty) => {
impl $crate::field::ArrowEnableVecForType for $t {}
};
}
#[doc(hidden)]
pub trait ArrowEnableVecForType {}
macro_rules! impl_numeric_type {
($physical_type:ty, $logical_type:ident) => {
impl ArrowField for $physical_type {
type Type = $physical_type;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::$logical_type
}
}
};
}
macro_rules! impl_numeric_type_full {
($physical_type:ty, $logical_type:ident) => {
impl_numeric_type!($physical_type, $logical_type);
arrow_enable_vec_for_type!($physical_type);
};
}
impl<T> ArrowField for Option<T>
where
T: ArrowField,
{
type Type = Option<<T as ArrowField>::Type>;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
<T as ArrowField>::data_type()
}
#[inline]
fn is_nullable() -> bool {
true
}
}
impl_numeric_type!(u8, UInt8);
impl_numeric_type_full!(u16, UInt16);
impl_numeric_type_full!(u32, UInt32);
impl_numeric_type_full!(u64, UInt64);
impl_numeric_type_full!(i8, Int8);
impl_numeric_type_full!(i16, Int16);
impl_numeric_type_full!(i32, Int32);
impl_numeric_type_full!(i64, Int64);
impl_numeric_type_full!(arrow2::types::f16, Float16);
impl_numeric_type_full!(f32, Float32);
impl_numeric_type_full!(f64, Float64);
pub struct I128<const PRECISION: usize, const SCALE: usize> {}
impl<const PRECISION: usize, const SCALE: usize> ArrowField for I128<PRECISION, SCALE> {
type Type = i128;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::Decimal(PRECISION, SCALE)
}
}
impl ArrowField for String {
type Type = String;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::Utf8
}
}
pub struct LargeString {}
impl ArrowField for LargeString {
type Type = String;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::LargeUtf8
}
}
impl ArrowField for bool {
type Type = Self;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::Boolean
}
}
impl ArrowField for NaiveDateTime {
type Type = Self;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::Timestamp(arrow2::datatypes::TimeUnit::Nanosecond, None)
}
}
impl ArrowField for NaiveDate {
type Type = Self;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::Date32
}
}
impl ArrowField for Buffer<u8> {
type Type = Self;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::Binary
}
}
impl ArrowField for Vec<u8> {
type Type = Self;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::Binary
}
}
pub struct LargeBinary {}
impl ArrowField for LargeBinary {
type Type = Vec<u8>;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::LargeBinary
}
}
pub struct FixedSizeBinary<const SIZE: usize> {}
impl<const SIZE: usize> ArrowField for FixedSizeBinary<SIZE> {
type Type = Vec<u8>;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::FixedSizeBinary(SIZE)
}
}
impl<T> ArrowField for Buffer<T>
where
T: ArrowField + NativeType + ArrowEnableVecForType,
{
type Type = Self;
#[inline]
fn data_type() -> DataType {
DataType::List(Box::new(<T as ArrowField>::field("item")))
}
}
impl<T> ArrowField for Vec<T>
where
T: ArrowField + ArrowEnableVecForType,
{
type Type = Vec<<T as ArrowField>::Type>;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::List(Box::new(<T as ArrowField>::field("item")))
}
}
pub struct LargeVec<T> {
d: std::marker::PhantomData<T>,
}
impl<T> ArrowField for LargeVec<T>
where
T: ArrowField + ArrowEnableVecForType,
{
type Type = Vec<<T as ArrowField>::Type>;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::LargeList(Box::new(<T as ArrowField>::field("item")))
}
}
pub struct FixedSizeVec<T, const SIZE: usize> {
d: std::marker::PhantomData<T>,
}
impl<T, const SIZE: usize> ArrowField for FixedSizeVec<T, SIZE>
where
T: ArrowField + ArrowEnableVecForType,
{
type Type = Vec<<T as ArrowField>::Type>;
#[inline]
fn data_type() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::FixedSizeList(Box::new(<T as ArrowField>::field("item")), SIZE)
}
}
arrow_enable_vec_for_type!(String);
arrow_enable_vec_for_type!(LargeString);
arrow_enable_vec_for_type!(bool);
arrow_enable_vec_for_type!(NaiveDateTime);
arrow_enable_vec_for_type!(NaiveDate);
arrow_enable_vec_for_type!(Vec<u8>);
arrow_enable_vec_for_type!(Buffer<u8>);
arrow_enable_vec_for_type!(LargeBinary);
impl<const SIZE: usize> ArrowEnableVecForType for FixedSizeBinary<SIZE> {}
impl<const PRECISION: usize, const SCALE: usize> ArrowEnableVecForType for I128<PRECISION, SCALE> {}
impl<T> ArrowEnableVecForType for Option<T> where T: ArrowField + ArrowEnableVecForType {}
impl<T> ArrowEnableVecForType for Vec<T> where T: ArrowField + ArrowEnableVecForType {}
impl<T> ArrowEnableVecForType for Buffer<T> where T: ArrowField + ArrowEnableVecForType {}
impl<T> ArrowEnableVecForType for LargeVec<T> where T: ArrowField + ArrowEnableVecForType {}
impl<T, const SIZE: usize> ArrowEnableVecForType for FixedSizeVec<T, SIZE> where
T: ArrowField + ArrowEnableVecForType
{
}