Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
[WIP] Add month and day temporal extractors
Browse files Browse the repository at this point in the history
  • Loading branch information
@VasanthakumarV
VasanthakumarV committed Sep 21, 2021
1 parent 7dedd02 commit fb7849f
Show file tree
Hide file tree
Showing 2 changed files with 132 additions and 98 deletions.
210 changes: 112 additions & 98 deletions src/compute/temporal.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,51 +75,33 @@ where
}
}

#[cfg(feature = "chrono-tz")]
#[cfg_attr(docsrs, doc(cfg(feature = "chrono-tz")))]
fn chrono_tz_hour(
array: &PrimitiveArray<i64>,
time_unit: TimeUnit,
timezone_str: &str,
) -> Result<PrimitiveArray<u32>> {
let timezone = parse_offset_tz(timezone_str)?;
Ok(extract_impl(array, time_unit, timezone, |x| x.hour()))
}

#[cfg(not(feature = "chrono-tz"))]
fn chrono_tz_hour(
_: &PrimitiveArray<i64>,
_: TimeUnit,
timezone_str: &str,
) -> Result<PrimitiveArray<u32>> {
Err(ArrowError::InvalidArgumentError(format!(
"timezone \"{}\" cannot be parsed (feature chrono-tz is not active)",
timezone_str
)))
}
macro_rules! chrono_tz {
($func:ident, $component:ident, $output:ty) => {
#[cfg(feature = "chrono-tz")]
#[cfg_attr(docsrs, doc(cfg(feature = "chrono-tz")))]
fn $func(
array: &PrimitiveArray<i64>,
time_unit: TimeUnit,
timezone_str: &str,
) -> Result<$output> {
let timezone = parse_offset_tz(timezone_str)?;
Ok(extract_impl(array, time_unit, timezone, |x| x.$component()))
}

#[cfg(feature = "chrono-tz")]
#[cfg_attr(docsrs, doc(cfg(feature = "chrono-tz")))]
fn chrono_tz_year(
array: &PrimitiveArray<i64>,
time_unit: TimeUnit,
timezone_str: &str,
) -> Result<PrimitiveArray<i32>> {
let timezone = parse_offset_tz(timezone_str)?;
Ok(extract_impl(array, time_unit, timezone, |x| x.year()))
#[cfg(not(feature = "chrono-tz"))]
fn $func(_: &PrimitiveArray<i64>, _: TimeUnit, timezone_str: &str) -> Result<$output> {
Err(ArrowError::InvalidArgumentError(format!(
"timezone \"{}\" cannot be parsed (feature chrono-tz is not active)",
timezone_str
)))
}
};
}

#[cfg(not(feature = "chrono-tz"))]
fn chrono_tz_year(
_: &PrimitiveArray<i64>,
_: TimeUnit,
timezone_str: &str,
) -> Result<PrimitiveArray<i32>> {
Err(ArrowError::InvalidArgumentError(format!(
"timezone \"{}\" cannot be parsed (feature chrono-tz is not active)",
timezone_str
)))
}
chrono_tz!(chrono_tz_hour, hour, PrimitiveArray<u32>);
chrono_tz!(chrono_tz_year, year, PrimitiveArray<i32>);
chrono_tz!(chrono_tz_month, month, PrimitiveArray<u32>);
chrono_tz!(chrono_tz_day, day, PrimitiveArray<u32>);

/// Extracts the hours of a temporal array as [`PrimitiveArray<u32>`].
/// Use [`can_hour`] to check if this operation is supported for the target [`DataType`].
Expand Down Expand Up @@ -246,63 +228,90 @@ pub fn can_hour(data_type: &DataType) -> bool {
)
}

macro_rules! date_like {
($component:ident, $array:ident, $data_type:path, $chrono_tz:ident) => {
match $array.data_type() {
DataType::Date32 => {
let array = $array
.as_any()
.downcast_ref::<PrimitiveArray<i32>>()
.unwrap();
Ok(unary(
array,
|x| date32_to_datetime(x).$component(),
$data_type,
))
}
DataType::Date64 => {
let array = $array
.as_any()
.downcast_ref::<PrimitiveArray<i64>>()
.unwrap();
Ok(unary(
array,
|x| date64_to_datetime(x).$component(),
$data_type,
))
}
DataType::Timestamp(time_unit, None) => {
let array = $array
.as_any()
.downcast_ref::<PrimitiveArray<i64>>()
.unwrap();
let op = match time_unit {
TimeUnit::Second => |x| timestamp_s_to_datetime(x).$component(),
TimeUnit::Millisecond => |x| timestamp_ms_to_datetime(x).$component(),
TimeUnit::Microsecond => |x| timestamp_us_to_datetime(x).$component(),
TimeUnit::Nanosecond => |x| timestamp_ns_to_datetime(x).$component(),
};
Ok(unary(array, op, $data_type))
}
DataType::Timestamp(time_unit, Some(timezone_str)) => {
let time_unit = *time_unit;
let timezone = parse_offset(timezone_str);

let array = $array.as_any().downcast_ref().unwrap();

if let Ok(timezone) = timezone {
Ok(extract_impl(array, time_unit, timezone, |x| x.$component()))
} else {
$chrono_tz(array, time_unit, timezone_str)
}
}
dt => Err(ArrowError::NotYetImplemented(format!(
"\"{}\" does not support type {:?}",
stringify!($component),
dt
))),
}
};
}

/// Extracts the years of a temporal array as [`PrimitiveArray<i32>`].
/// Use [`can_year`] to check if this operation is supported for the target [`DataType`].
pub fn year(array: &dyn Array) -> Result<PrimitiveArray<i32>> {
let final_data_type = DataType::Int32;
match array.data_type() {
DataType::Date32 => {
let array = array
.as_any()
.downcast_ref::<PrimitiveArray<i32>>()
.unwrap();
Ok(unary(
array,
|x| date32_to_datetime(x).year(),
final_data_type,
))
}
DataType::Date64 => {
let array = array
.as_any()
.downcast_ref::<PrimitiveArray<i64>>()
.unwrap();
Ok(unary(
array,
|x| date64_to_datetime(x).year(),
final_data_type,
))
}
DataType::Timestamp(time_unit, None) => {
let array = array
.as_any()
.downcast_ref::<PrimitiveArray<i64>>()
.unwrap();
let op = match time_unit {
TimeUnit::Second => |x| timestamp_s_to_datetime(x).year(),
TimeUnit::Millisecond => |x| timestamp_ms_to_datetime(x).year(),
TimeUnit::Microsecond => |x| timestamp_us_to_datetime(x).year(),
TimeUnit::Nanosecond => |x| timestamp_ns_to_datetime(x).year(),
};
Ok(unary(array, op, final_data_type))
}
DataType::Timestamp(time_unit, Some(timezone_str)) => {
let time_unit = *time_unit;
let timezone = parse_offset(timezone_str);
date_like!(year, array, DataType::Int32, chrono_tz_year)
}

let array = array.as_any().downcast_ref().unwrap();
/// Extracts the months of a temporal array as [`PrimitiveArray<u32>`].
/// Use [`can_month`] to check if this operation is supported for the target [`DataType`].
pub fn month(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
date_like!(month, array, DataType::UInt32, chrono_tz_month)
}

if let Ok(timezone) = timezone {
Ok(extract_impl(array, time_unit, timezone, |x| x.year()))
} else {
chrono_tz_year(array, time_unit, timezone_str)
}
}
dt => Err(ArrowError::NotYetImplemented(format!(
"\"year\" does not support type {:?}",
dt
))),
}
/// Extracts the days of a temporal array as [`PrimitiveArray<u32>`].
/// Use [`can_month`] to check if this operation is supported for the target [`DataType`].
pub fn day(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
date_like!(day, array, DataType::UInt32, chrono_tz_day)
}

macro_rules! can_date_like {
($dt:expr) => {
matches!(
$dt,
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
)
};
}

/// Checks if an array of type `datatype` can perform year operation
Expand All @@ -319,8 +328,13 @@ pub fn year(array: &dyn Array) -> Result<PrimitiveArray<i32>> {
/// assert_eq!(can_year(&data_type), false);
/// ```
pub fn can_year(data_type: &DataType) -> bool {
matches!(
data_type,
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
)
can_date_like!(data_type)
}

pub fn can_month(data_type: &DataType) -> bool {
can_date_like!(data_type)
}

pub fn can_day(data_type: &DataType) -> bool {
can_date_like!(data_type)
}
20 changes: 20 additions & 0 deletions tests/it/compute/temporal.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,26 @@ fn date64_year() {
assert_eq!(result, expected);
}

#[test]
fn date64_month() {
let array = Int64Array::from(&[Some(1514764800000), None]).to(DataType::Date64);
let result = month(&array).unwrap();

let expected = UInt32Array::from(&[Some(1), None]);

assert_eq!(result, expected);
}

#[test]
fn date64_day() {
let array = Int64Array::from(&[Some(1614764800000), None]).to(DataType::Date64);
let result = day(&array).unwrap();

let expected = UInt32Array::from(&[Some(3), None]);

assert_eq!(result, expected);
}

#[test]
fn naive_timestamp_date32_year() {
let array = Int32Array::from(&[Some(15147), None]).to(DataType::Date32);
Expand Down

0 comments on commit fb7849f

Please sign in to comment.