Skip to content

Commit

Permalink
feat(compute): Support doy (day of year) for temporal
Browse files Browse the repository at this point in the history
  • Loading branch information
ovr committed Jul 14, 2022
1 parent 4444cb7 commit 92b1aca
Showing 1 changed file with 50 additions and 1 deletion.
51 changes: 50 additions & 1 deletion arrow/src/compute/kernels/temporal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -338,7 +338,7 @@ where
scratch
)
}
dt => return_compute_error_with!("weekday does not support", dt),
dt => return_compute_error_with!("weekday0 does not support", dt),
}

Ok(b.finish())
Expand Down Expand Up @@ -372,6 +372,35 @@ where
Ok(b.finish())
}

/// Extracts the day of year of a given temporal array as an array of integers
/// The day of year that ranges from 1 to 366
pub fn doy<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
where
T: ArrowTemporalType + ArrowNumericType,
i64: std::convert::From<T::Native>,
{
let mut b = Int32Builder::new(array.len());
match array.data_type() {
&DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, None) => {
extract_component_from_array!(array, b, ordinal, value_as_datetime)
}
&DataType::Timestamp(_, Some(ref tz)) => {
let mut scratch = Parsed::new();
extract_component_from_array!(
array,
b,
ordinal,
value_as_datetime_with_tz,
tz,
scratch
)
}
dt => return_compute_error_with!("doy does not support", dt),
}

Ok(b.finish())
}

/// Extracts the minutes of a given temporal array as an array of integers
pub fn minute<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
where
Expand Down Expand Up @@ -681,6 +710,26 @@ mod tests {
assert_eq!(1, b.value(2));
}

#[test]
fn test_temporal_array_date64_doy() {
//1483228800000 -> 2017-01-01 (Sunday)
//1514764800000 -> 2018-01-01
//1550636625000 -> 2019-02-20
let a: PrimitiveArray<Date64Type> = vec![
Some(1483228800000),
Some(1514764800000),
None,
Some(1550636625000),
]
.into();

let b = doy(&a).unwrap();
assert_eq!(1, b.value(0));
assert_eq!(1, b.value(1));
assert!(!b.is_valid(2));
assert_eq!(51, b.value(3));
}

#[test]
fn test_temporal_array_timestamp_micro_year() {
let a: TimestampMicrosecondArray =
Expand Down

0 comments on commit 92b1aca

Please sign in to comment.