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

Commit

Permalink
Add month and day temporal extractors
Browse files Browse the repository at this point in the history
  • Loading branch information
@VasanthakumarV
VasanthakumarV committed Sep 22, 2021
1 parent 7dedd02 commit a5cfec4
Show file tree
Hide file tree
Showing 2 changed files with 116 additions and 34 deletions.
130 changes: 96 additions & 34 deletions src/compute/temporal.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,21 +100,31 @@ fn chrono_tz_hour(

#[cfg(feature = "chrono-tz")]
#[cfg_attr(docsrs, doc(cfg(feature = "chrono-tz")))]
fn chrono_tz_year(
fn chrono_tz<F, O>(
array: &PrimitiveArray<i64>,
time_unit: TimeUnit,
timezone_str: &str,
) -> Result<PrimitiveArray<i32>> {
op: F,
) -> Result<PrimitiveArray<O>>
where
O: NativeType,
F: Fn(chrono::DateTime<chrono_tz::Tz>) -> O,
{
let timezone = parse_offset_tz(timezone_str)?;
Ok(extract_impl(array, time_unit, timezone, |x| x.year()))
Ok(extract_impl(array, time_unit, timezone, op))
}

#[cfg(not(feature = "chrono-tz"))]
fn chrono_tz_year(
fn chrono_tz<F, O>(
_: &PrimitiveArray<i64>,
_: TimeUnit,
timezone_str: &str,
) -> Result<PrimitiveArray<i32>> {
_: F,
) -> Result<PrimitiveArray<O>>
where
O: NativeType,
F: Fn(chrono::DateTime<chrono_tz::Tz>) -> O,
{
Err(ArrowError::InvalidArgumentError(format!(
"timezone \"{}\" cannot be parsed (feature chrono-tz is not active)",
timezone_str
Expand Down Expand Up @@ -249,54 +259,92 @@ pub fn can_hour(data_type: &DataType) -> bool {
/// 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;
if let DataType::Timestamp(time_unit, Some(timezone_str)) = array.data_type() {
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.year()))
} else {
chrono_tz(array, time_unit, timezone_str, |x| x.year())
}
} else {
date_like(array, DataType::Int32, |x| x.year())
}
}

/// 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>> {
if let DataType::Timestamp(time_unit, Some(timezone_str)) = array.data_type() {
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.month()))
} else {
chrono_tz(array, time_unit, timezone_str, |x| x.month())
}
} else {
date_like(array, DataType::UInt32, |x| x.month())
}
}

/// Extracts the days of a temporal array as [`PrimitiveArray<u32>`].
/// Use [`can_day`] to check if this operation is supported for the target [`DataType`].
pub fn day(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
if let DataType::Timestamp(time_unit, Some(timezone_str)) = array.data_type() {
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.day()))
} else {
chrono_tz(array, time_unit, timezone_str, |x| x.day())
}
} else {
date_like(array, DataType::UInt32, |x| x.day())
}
}

pub fn date_like<F, O>(array: &dyn Array, data_type: DataType, op: F) -> Result<PrimitiveArray<O>>
where
O: NativeType,
F: Fn(chrono::NaiveDateTime) -> O,
{
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,
))
Ok(unary(array, |x| op(date32_to_datetime(x)), 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,
))
Ok(unary(array, |x| op(date64_to_datetime(x)), 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(),
let func = match time_unit {
TimeUnit::Second => timestamp_s_to_datetime,
TimeUnit::Millisecond => timestamp_ms_to_datetime,
TimeUnit::Microsecond => timestamp_us_to_datetime,
TimeUnit::Nanosecond => timestamp_ns_to_datetime,
};
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);

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

if let Ok(timezone) = timezone {
Ok(extract_impl(array, time_unit, timezone, |x| x.year()))
} else {
chrono_tz_year(array, time_unit, timezone_str)
}
Ok(unary(array, |x| op(func(x)), data_type))
}
dt => Err(ArrowError::NotYetImplemented(format!(
"\"year\" does not support type {:?}",
Expand Down Expand Up @@ -324,3 +372,17 @@ pub fn can_year(data_type: &DataType) -> bool {
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
)
}

pub fn can_month(data_type: &DataType) -> bool {
matches!(
data_type,
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
)
}

pub fn can_day(data_type: &DataType) -> bool {
matches!(
data_type,
DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
)
}
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 @@ -78,6 +78,26 @@ fn naive_timestamp_micro_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 timestamp_micro_hour() {
let array = Int64Array::from(&[Some(1621877130000000), None]).to(DataType::Timestamp(
Expand Down

0 comments on commit a5cfec4

Please sign in to comment.