use arrow2::array::*;
use arrow2::compute::temporal::*;
use arrow2::datatypes::*;
macro_rules! temporal_test {
($func:ident, $extract:ident, $data_types:path) => {
#[test]
fn $func() {
for data_type in $data_types() {
let data = TestData::data(&data_type);
let result = $extract(&*data.input).unwrap();
assert_eq!(
result,
data.$extract.unwrap(),
"\"{}\" failed on type: {:?}",
stringify!($extract),
data_type
);
}
}
};
}
temporal_test!(temporal_hour, hour, TestData::available_time_like_types);
temporal_test!(temporal_minute, minute, TestData::available_time_like_types);
temporal_test!(temporal_second, second, TestData::available_time_like_types);
temporal_test!(
temporal_nanosecond,
nanosecond,
TestData::available_time_like_types
);
temporal_test!(temporal_year, year, TestData::available_date_like_types);
temporal_test!(temporal_month, month, TestData::available_date_like_types);
temporal_test!(temporal_day, day, TestData::available_date_like_types);
temporal_test!(
temporal_weekday,
weekday,
TestData::available_date_like_types
);
temporal_test!(
temporal_iso_week,
iso_week,
TestData::available_date_like_types
);
struct TestData {
input: Box<dyn Array>,
year: Option<Int32Array>,
month: Option<UInt32Array>,
day: Option<UInt32Array>,
weekday: Option<UInt32Array>,
iso_week: Option<UInt32Array>,
hour: Option<UInt32Array>,
minute: Option<UInt32Array>,
second: Option<UInt32Array>,
nanosecond: Option<UInt32Array>,
}
impl TestData {
fn data(data_type: &DataType) -> TestData {
match data_type {
DataType::Date64 => TestData {
input: Box::new(
Int64Array::from(&[Some(1514764800000), None, Some(1550636625000)])
.to(data_type.clone()),
),
year: Some(Int32Array::from(&[Some(2018), None, Some(2019)])),
month: Some(UInt32Array::from(&[Some(1), None, Some(2)])),
day: Some(UInt32Array::from(&[Some(1), None, Some(20)])),
weekday: Some(UInt32Array::from(&[Some(1), None, Some(3)])),
iso_week: Some(UInt32Array::from(&[Some(1), None, Some(8)])),
hour: Some(UInt32Array::from(&[Some(0), None, Some(4)])),
minute: Some(UInt32Array::from(&[Some(0), None, Some(23)])),
second: Some(UInt32Array::from(&[Some(0), None, Some(45)])),
nanosecond: Some(UInt32Array::from(&[Some(0), None, Some(0)])),
},
DataType::Date32 => TestData {
input: Box::new(Int32Array::from(&[Some(15147), None]).to(data_type.clone())),
year: Some(Int32Array::from(&[Some(2011), None])),
month: Some(UInt32Array::from(&[Some(6), None])),
day: Some(UInt32Array::from(&[Some(22), None])),
weekday: Some(UInt32Array::from(&[Some(3), None])),
iso_week: Some(UInt32Array::from(&[Some(25), None])),
hour: Some(UInt32Array::from(&[Some(0), None])),
minute: Some(UInt32Array::from(&[Some(0), None])),
second: Some(UInt32Array::from(&[Some(0), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
DataType::Time32(TimeUnit::Second) => TestData {
input: Box::new(Int32Array::from(&[Some(37800), None]).to(data_type.clone())),
year: None,
month: None,
day: None,
weekday: None,
iso_week: None,
hour: Some(UInt32Array::from(&[Some(10), None])),
minute: Some(UInt32Array::from(&[Some(30), None])),
second: Some(UInt32Array::from(&[Some(0), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
DataType::Time64(TimeUnit::Microsecond) => TestData {
input: Box::new(Int64Array::from(&[Some(378000000), None]).to(data_type.clone())),
year: None,
month: None,
day: None,
weekday: None,
iso_week: None,
hour: Some(UInt32Array::from(&[Some(0), None])),
minute: Some(UInt32Array::from(&[Some(6), None])),
second: Some(UInt32Array::from(&[Some(18), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
DataType::Time64(TimeUnit::Nanosecond) => TestData {
input: Box::new(
Int64Array::from(&[Some(378000000100), None]).to(data_type.clone()),
),
year: None,
month: None,
day: None,
weekday: None,
iso_week: None,
hour: Some(UInt32Array::from(&[Some(0), None])),
minute: Some(UInt32Array::from(&[Some(6), None])),
second: Some(UInt32Array::from(&[Some(18), None])),
nanosecond: Some(UInt32Array::from(&[Some(100), None])),
},
DataType::Timestamp(TimeUnit::Microsecond, None) => TestData {
input: Box::new(
Int64Array::from(&[Some(1612025847000000), None, Some(68216970000000)])
.to(data_type.clone()),
),
year: Some(Int32Array::from(&[Some(2021), None, Some(1972)])),
month: Some(UInt32Array::from(&[Some(1), None, Some(2)])),
day: Some(UInt32Array::from(&[Some(30), None, Some(29)])),
weekday: Some(UInt32Array::from(&[Some(6), None, Some(2)])),
iso_week: Some(UInt32Array::from(&[Some(4), None, Some(9)])),
hour: Some(UInt32Array::from(&[Some(16), None, Some(13)])),
minute: Some(UInt32Array::from(&[Some(57), None, Some(9)])),
second: Some(UInt32Array::from(&[Some(27), None, Some(30)])),
nanosecond: Some(UInt32Array::from(&[Some(0), None, Some(0)])),
},
_ => unreachable!(),
}
}
fn available_time_like_types() -> Vec<DataType> {
vec![
DataType::Date32,
DataType::Date64,
DataType::Time32(TimeUnit::Second),
DataType::Time64(TimeUnit::Microsecond),
DataType::Time64(TimeUnit::Nanosecond),
DataType::Timestamp(TimeUnit::Microsecond, None),
]
}
fn available_date_like_types() -> Vec<DataType> {
vec![
DataType::Date32,
DataType::Date64,
DataType::Timestamp(TimeUnit::Microsecond, None),
]
}
}
macro_rules! temporal_tz_test {
($func:ident, $extract:ident) => {
#[cfg(feature = "chrono-tz")]
#[test]
fn $func() {
let test_data = test_data_tz();
for data in test_data {
let result = $extract(&*data.input).unwrap();
assert_eq!(result, data.$extract.unwrap());
}
}
};
}
temporal_tz_test!(temporal_tz_hour, hour);
temporal_tz_test!(temporal_tz_minute, minute);
temporal_tz_test!(temporal_tz_second, second);
temporal_tz_test!(temporal_tz_nanosecond, nanosecond);
temporal_tz_test!(temporal_tz_year, year);
temporal_tz_test!(temporal_tz_month, month);
temporal_tz_test!(temporal_tz_day, day);
temporal_tz_test!(temporal_tz_weekday, weekday);
temporal_tz_test!(temporal_tz_iso_week, iso_week);
fn test_data_tz() -> Vec<TestData> {
vec![
TestData {
input: Box::new(
Int64Array::from(&[Some(1621877130000000), None]).to(DataType::Timestamp(
TimeUnit::Microsecond,
Some("GMT".to_string()),
)),
),
year: Some(Int32Array::from(&[Some(2021), None])),
month: Some(UInt32Array::from(&[Some(5), None])),
day: Some(UInt32Array::from(&[Some(24), None])),
weekday: Some(UInt32Array::from(&[Some(1), None])),
iso_week: Some(UInt32Array::from(&[Some(21), None])),
hour: Some(UInt32Array::from(&[Some(17), None])),
minute: Some(UInt32Array::from(&[Some(25), None])),
second: Some(UInt32Array::from(&[Some(30), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
TestData {
input: Box::new(Int64Array::from(&[Some(1621877130000000), None]).to(
DataType::Timestamp(TimeUnit::Microsecond, Some("+01:00".to_string())),
)),
year: Some(Int32Array::from(&[Some(2021), None])),
month: Some(UInt32Array::from(&[Some(5), None])),
day: Some(UInt32Array::from(&[Some(24), None])),
weekday: Some(UInt32Array::from(&[Some(1), None])),
iso_week: Some(UInt32Array::from(&[Some(21), None])),
hour: Some(UInt32Array::from(&[Some(18), None])),
minute: Some(UInt32Array::from(&[Some(25), None])),
second: Some(UInt32Array::from(&[Some(30), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
TestData {
input: Box::new(Int64Array::from(&[Some(1621877130000000), None]).to(
DataType::Timestamp(TimeUnit::Microsecond, Some("Europe/Lisbon".to_string())),
)),
year: Some(Int32Array::from(&[Some(2021), None])),
month: Some(UInt32Array::from(&[Some(5), None])),
day: Some(UInt32Array::from(&[Some(24), None])),
weekday: Some(UInt32Array::from(&[Some(1), None])),
iso_week: Some(UInt32Array::from(&[Some(21), None])),
hour: Some(UInt32Array::from(&[Some(18), None])),
minute: Some(UInt32Array::from(&[Some(25), None])),
second: Some(UInt32Array::from(&[Some(30), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
TestData {
input: Box::new(
Int64Array::from(&[Some(1585440000), None]).to(DataType::Timestamp(
TimeUnit::Second,
Some("Europe/Lisbon".to_string()),
)),
),
year: Some(Int32Array::from(&[Some(2020), None])),
month: Some(UInt32Array::from(&[Some(3), None])),
day: Some(UInt32Array::from(&[Some(29), None])),
weekday: Some(UInt32Array::from(&[Some(7), None])),
iso_week: Some(UInt32Array::from(&[Some(13), None])),
hour: Some(UInt32Array::from(&[Some(0), None])),
minute: Some(UInt32Array::from(&[Some(0), None])),
second: Some(UInt32Array::from(&[Some(0), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
TestData {
input: Box::new(
Int64Array::from(&[Some(1585443600), None]).to(DataType::Timestamp(
TimeUnit::Second,
Some("Europe/Lisbon".to_string()),
)),
),
year: Some(Int32Array::from(&[Some(2020), None])),
month: Some(UInt32Array::from(&[Some(3), None])),
day: Some(UInt32Array::from(&[Some(29), None])),
weekday: Some(UInt32Array::from(&[Some(7), None])),
iso_week: Some(UInt32Array::from(&[Some(13), None])),
hour: Some(UInt32Array::from(&[Some(2), None])),
minute: Some(UInt32Array::from(&[Some(0), None])),
second: Some(UInt32Array::from(&[Some(0), None])),
nanosecond: Some(UInt32Array::from(&[Some(0), None])),
},
]
}
#[test]
fn consistency_hour() {
consistency_check(can_hour, hour);
}
#[test]
fn consistency_minute() {
consistency_check(can_minute, minute);
}
#[test]
fn consistency_second() {
consistency_check(can_second, second);
}
#[test]
fn consistency_nanosecond() {
consistency_check(can_nanosecond, nanosecond);
}
#[test]
fn consistency_year() {
consistency_check(can_year, year);
}
#[test]
fn consistency_month() {
consistency_check(can_month, month);
}
#[test]
fn consistency_day() {
consistency_check(can_day, day);
}
#[test]
fn consistency_weekday() {
consistency_check(can_weekday, weekday);
}
#[test]
fn consistency_iso_week() {
consistency_check(can_iso_week, iso_week);
}
fn consistency_check<O: arrow2::types::NativeType>(
can_extract: fn(&DataType) -> bool,
extract: fn(&dyn Array) -> arrow2::error::Result<PrimitiveArray<O>>,
) {
use arrow2::datatypes::DataType::*;
let datatypes = vec![
Null,
Boolean,
UInt8,
UInt16,
UInt32,
UInt64,
Int8,
Int16,
Int32,
Int64,
Float32,
Float64,
Timestamp(TimeUnit::Second, None),
Timestamp(TimeUnit::Millisecond, None),
Timestamp(TimeUnit::Microsecond, None),
Timestamp(TimeUnit::Nanosecond, None),
Timestamp(TimeUnit::Nanosecond, Some("+00:00".to_string())),
Time64(TimeUnit::Microsecond),
Time64(TimeUnit::Nanosecond),
Date32,
Time32(TimeUnit::Second),
Time32(TimeUnit::Millisecond),
Date64,
Utf8,
LargeUtf8,
Binary,
LargeBinary,
Duration(TimeUnit::Second),
Duration(TimeUnit::Millisecond),
Duration(TimeUnit::Microsecond),
Duration(TimeUnit::Nanosecond),
];
datatypes.into_iter().for_each(|d1| {
let array = new_null_array(d1.clone(), 10);
if can_extract(&d1) {
assert!(extract(array.as_ref()).is_ok());
} else {
assert!(extract(array.as_ref()).is_err());
}
});
}