use std::ops::{Deref, Div};
use arrow::temporal_conversions::{MICROSECONDS_IN_DAY, MILLISECONDS_IN_DAY, NANOSECONDS_IN_DAY};
use polars_core::prelude::arity::unary_elementwise_values;
use polars_core::prelude::*;
use crate::chunkedarray::*;
pub trait AsSeries {
fn as_series(&self) -> &Series;
}
impl AsSeries for Series {
fn as_series(&self) -> &Series {
self
}
}
pub trait TemporalMethods: AsSeries {
fn hour(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.hour()),
#[cfg(feature = "dtype-time")]
DataType::Time => s.time().map(|ca| ca.hour()),
dt => polars_bail!(opq = hour, dt),
}
}
fn minute(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.minute()),
#[cfg(feature = "dtype-time")]
DataType::Time => s.time().map(|ca| ca.minute()),
dt => polars_bail!(opq = minute, dt),
}
}
fn second(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.second()),
#[cfg(feature = "dtype-time")]
DataType::Time => s.time().map(|ca| ca.second()),
dt => polars_bail!(opq = second, dt),
}
}
fn nanosecond(&self) -> PolarsResult<Int32Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.nanosecond()),
#[cfg(feature = "dtype-time")]
DataType::Time => s.time().map(|ca| ca.nanosecond()),
dt => polars_bail!(opq = nanosecond, dt),
}
}
fn day(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.day()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.day()),
dt => polars_bail!(opq = day, dt),
}
}
fn weekday(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| {
unary_elementwise_values(ca, |t| (((t - 4) % 7 + 7) % 7 + 1) as i8)
}),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(time_unit, time_zone) => s.datetime().map(|ca| {
match time_zone.as_deref() {
Some("UTC") | None => {
let divisor = match time_unit {
TimeUnit::Milliseconds => MILLISECONDS_IN_DAY,
TimeUnit::Microseconds => MICROSECONDS_IN_DAY,
TimeUnit::Nanoseconds => NANOSECONDS_IN_DAY,
};
unary_elementwise_values(ca, |t| {
let t = t / divisor - ((t < 0 && t % divisor != 0) as i64);
(((t - 4) % 7 + 7) % 7 + 1) as i8
})
},
_ => ca.weekday(),
}
}),
dt => polars_bail!(opq = weekday, dt),
}
}
fn week(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.week()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.week()),
dt => polars_bail!(opq = week, dt),
}
}
fn ordinal_day(&self) -> PolarsResult<Int16Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.ordinal()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.ordinal()),
dt => polars_bail!(opq = ordinal_day, dt),
}
}
fn millennium(&self) -> PolarsResult<Int32Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| (ca.year() - 1i32).div(1000f64) + 1),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| (ca.year() - 1i32).div(1000f64) + 1),
dt => polars_bail!(opq = century, dt),
}
}
fn century(&self) -> PolarsResult<Int32Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| (ca.year() - 1i32).div(100f64) + 1),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| (ca.year() - 1i32).div(100f64) + 1),
dt => polars_bail!(opq = century, dt),
}
}
fn year(&self) -> PolarsResult<Int32Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.year()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.year()),
dt => polars_bail!(opq = year, dt),
}
}
fn iso_year(&self) -> PolarsResult<Int32Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.iso_year()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.iso_year()),
dt => polars_bail!(opq = iso_year, dt),
}
}
fn ordinal_year(&self) -> PolarsResult<Int32Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.year()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.year()),
dt => polars_bail!(opq = ordinal_year, dt),
}
}
fn is_leap_year(&self) -> PolarsResult<BooleanChunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.is_leap_year()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.is_leap_year()),
dt => polars_bail!(opq = is_leap_year, dt),
}
}
fn quarter(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.quarter()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.quarter()),
dt => polars_bail!(opq = quarter, dt),
}
}
fn month(&self) -> PolarsResult<Int8Chunked> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-date")]
DataType::Date => s.date().map(|ca| ca.month()),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => s.datetime().map(|ca| ca.month()),
dt => polars_bail!(opq = month, dt),
}
}
fn to_string(&self, format: &str) -> PolarsResult<Series> {
let s = self.as_series();
match s.dtype() {
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(_, _) => {
let format = get_strftime_format(format, s.dtype())?;
s.datetime()
.map(|ca| Ok(ca.to_string(format.as_str())?.into_series()))?
},
#[cfg(feature = "dtype-date")]
DataType::Date => {
let format = get_strftime_format(format, s.dtype())?;
s.date()
.map(|ca| Ok(ca.to_string(format.as_str())?.into_series()))?
},
#[cfg(feature = "dtype-time")]
DataType::Time => {
let format = get_strftime_format(format, s.dtype())?;
s.time()
.map(|ca| ca.to_string(format.as_str()).into_series())
},
#[cfg(feature = "dtype-duration")]
DataType::Duration(_) => s
.duration()
.map(|ca| Ok(ca.to_string(format)?.into_series()))?,
dt => polars_bail!(opq = to_string, dt),
}
}
fn strftime(&self, format: &str) -> PolarsResult<Series> {
self.to_string(format)
}
#[cfg(feature = "temporal")]
fn timestamp(&self, tu: TimeUnit) -> PolarsResult<Int64Chunked> {
let s = self.as_series();
if matches!(s.dtype(), DataType::Time | DataType::Duration(_)) {
polars_bail!(opq = timestamp, s.dtype());
} else {
s.cast(&DataType::Datetime(tu, None))
.map(|s| s.datetime().unwrap().deref().clone())
}
}
}
impl<T: ?Sized + AsSeries> TemporalMethods for T {}