Skip to content

Commit

Permalink
jiff
Browse files Browse the repository at this point in the history
  • Loading branch information
ijl committed Oct 7, 2024
1 parent 4d3ff2c commit b1d487a
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 62 deletions.
32 changes: 7 additions & 25 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -59,12 +59,12 @@ strict_provenance = []
arrayvec = { version = "0.7", default-features = false, features = ["std", "serde"] }
associative-cache = { version = "2", default-features = false }
bytecount = { version = "^0.6.7", default-features = false, features = ["runtime-dispatch-simd"] }
chrono = { version = "=0.4.34", default-features = false }
compact_str = { version = "0.8", default-features = false, features = ["serde"] }
encoding_rs = { version = "0.8", default-features = false }
half = { version = "2", default-features = false, features = ["std"] }
itoa = { version = "1", default-features = false }
itoap = { version = "1", features = ["std", "simd"] }
itoap = { version = "1", default-features = false, features = ["std", "simd"] }
jiff = { version = "^0.1", default-features = false, features = ["alloc"] }
once_cell = { version = "1", default-features = false, features = ["alloc", "race"] }
pyo3-ffi = { path = "include/pyo3/pyo3-ffi", default-features = false, features = ["extension-module"]}
ryu = { version = "1", default-features = false }
Expand Down
75 changes: 40 additions & 35 deletions src/serialize/per_type/numpy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,9 @@ use crate::serialize::per_type::{
};
use crate::serialize::serializer::PyObjectSerializer;
use crate::typeref::{load_numpy_types, ARRAY_STRUCT_STR, DESCR_STR, DTYPE_STR, NUMPY_TYPES};
use chrono::{Datelike, NaiveDate, NaiveDateTime, Timelike};
use core::ffi::{c_char, c_int, c_void};
use jiff::civil::DateTime;
use jiff::Timestamp;
use pyo3_ffi::*;
use serde::ser::{self, Serialize, SerializeSeq, Serializer};
use std::fmt;
Expand Down Expand Up @@ -1203,6 +1204,19 @@ impl NumpyDateTimeError {
}
}

macro_rules! to_jiff_datetime {
($timestamp:expr, $self:expr, $val:expr) => {
Ok(
($timestamp.map_err(|_| NumpyDateTimeError::Unrepresentable {
unit: *$self,
val: $val,
})?)
.to_zoned(jiff::tz::TimeZone::UTC)
.datetime(),
)
};
}

impl NumpyDatetimeUnit {
/// Create a `NumpyDatetimeUnit` from a pointer to a Python object holding a
/// numpy array.
Expand Down Expand Up @@ -1255,56 +1269,44 @@ impl NumpyDatetimeUnit {
#[cfg_attr(feature = "optimize", optimize(size))]
fn datetime(&self, val: i64, opts: Opt) -> Result<NumpyDatetime64Repr, NumpyDateTimeError> {
match self {
Self::Years => Ok(NaiveDate::from_ymd_opt(
Self::Years => Ok(DateTime::new(
(val + 1970)
.try_into()
.map_err(|_| NumpyDateTimeError::Unrepresentable { unit: *self, val })?,
1,
1,
0,
0,
0,
0,
)
.unwrap()
.and_hms_opt(0, 0, 0)
.unwrap()),
Self::Months => Ok(NaiveDate::from_ymd_opt(
Self::Months => Ok(DateTime::new(
(val / 12 + 1970)
.try_into()
.map_err(|_| NumpyDateTimeError::Unrepresentable { unit: *self, val })?,
(val % 12 + 1)
.try_into()
.map_err(|_| NumpyDateTimeError::Unrepresentable { unit: *self, val })?,
1,
0,
0,
0,
0,
)
.unwrap()
.and_hms_opt(0, 0, 0)
.unwrap()),
Self::Weeks => {
Ok(NaiveDateTime::from_timestamp_opt(val * 7 * 24 * 60 * 60, 0).unwrap())
to_jiff_datetime!(Timestamp::from_second(val * 7 * 24 * 60 * 60), self, val)
}
Self::Days => to_jiff_datetime!(Timestamp::from_second(val * 24 * 60 * 60), self, val),
Self::Hours => to_jiff_datetime!(Timestamp::from_second(val * 60 * 60), self, val),
Self::Minutes => to_jiff_datetime!(Timestamp::from_second(val * 60), self, val),
Self::Seconds => to_jiff_datetime!(Timestamp::from_second(val), self, val),
Self::Milliseconds => to_jiff_datetime!(Timestamp::from_millisecond(val), self, val),
Self::Microseconds => to_jiff_datetime!(Timestamp::from_microsecond(val), self, val),
Self::Nanoseconds => {
to_jiff_datetime!(Timestamp::from_nanosecond(val as i128), self, val)
}
Self::Days => Ok(NaiveDateTime::from_timestamp_opt(val * 24 * 60 * 60, 0).unwrap()),
Self::Hours => Ok(NaiveDateTime::from_timestamp_opt(val * 60 * 60, 0).unwrap()),
Self::Minutes => Ok(NaiveDateTime::from_timestamp_opt(val * 60, 0).unwrap()),
Self::Seconds => Ok(NaiveDateTime::from_timestamp_opt(val, 0).unwrap()),
Self::Milliseconds => Ok(NaiveDateTime::from_timestamp_opt(
val / 1_000,
(val % 1_000 * 1_000_000)
.try_into()
.map_err(|_| NumpyDateTimeError::Unrepresentable { unit: *self, val })?,
)
.unwrap()),
Self::Microseconds => Ok(NaiveDateTime::from_timestamp_opt(
val / 1_000_000,
(val % 1_000_000 * 1_000)
.try_into()
.map_err(|_| NumpyDateTimeError::Unrepresentable { unit: *self, val })?,
)
.unwrap()),
Self::Nanoseconds => Ok(NaiveDateTime::from_timestamp_opt(
val / 1_000_000_000,
(val % 1_000_000_000)
.try_into()
.map_err(|_| NumpyDateTimeError::Unrepresentable { unit: *self, val })?,
)
.unwrap()),
_ => Err(NumpyDateTimeError::UnsupportedUnit(*self)),
}
.map(|dt| NumpyDatetime64Repr { dt, opts })
Expand Down Expand Up @@ -1357,7 +1359,7 @@ macro_rules! forward_inner {
}

struct NumpyDatetime64Repr {
dt: NaiveDateTime,
dt: DateTime,
opts: Opt,
}

Expand All @@ -1368,7 +1370,10 @@ impl DateTimeLike for NumpyDatetime64Repr {
forward_inner!(hour, u8);
forward_inner!(minute, u8);
forward_inner!(second, u8);
forward_inner!(nanosecond, u32);

fn nanosecond(&self) -> u32 {
self.dt.subsec_nanosecond() as u32
}

fn microsecond(&self) -> u32 {
self.nanosecond() / 1_000
Expand Down

0 comments on commit b1d487a

Please sign in to comment.