nd_arrow_array/compute/
time.rs

1//! CF time coordinates and time labels for resampling.
2//!
3//! [`decode_cf_time`] turns numbers with `units = "<unit> since <reference>"`
4//! into an Arrow `Timestamp(Microsecond)` without time zone, like xarray. The
5//! microsecond unit covers ±292,000 years. Only the `standard`, `gregorian` and
6//! `proleptic_gregorian` calendars are supported, all as proleptic Gregorian.
7
8use std::sync::Arc;
9
10use arrow::array::{ArrayRef, AsArray, Float64Array, TimestampMicrosecondArray};
11use arrow::compute::cast as arrow_cast;
12pub use arrow::compute::kernels::temporal::DatePart;
13use arrow::compute::kernels::temporal::date_part;
14use arrow::datatypes::{DataType, Float64Type, TimeUnit, TimestampMicrosecondType};
15
16use crate::compute::cf::decode_cf;
17use crate::error::{Result, nd_err};
18use crate::{Attributes, NdArrowArray};
19
20/// CF attribute: the time unit and reference, for example `days since 1950-01-01`.
21pub const UNITS: &str = "units";
22/// CF attribute: the calendar of the time values.
23pub const CALENDAR: &str = "calendar";
24
25pub(crate) const MICROS_PER_SECOND: i64 = 1_000_000;
26pub(crate) const MICROS_PER_DAY: i64 = 86_400 * MICROS_PER_SECOND;
27
28/// The decoded `Timestamp` type.
29pub fn timestamp_type() -> DataType {
30    DataType::Timestamp(TimeUnit::Microsecond, None)
31}
32
33/// A parsed CF time unit: microseconds per step and the reference instant.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct CfTimeUnits {
36    /// Microseconds per unit step.
37    pub step_micros: i64,
38    /// The reference instant, in microseconds since 1970-01-01.
39    pub reference_micros: i64,
40}
41
42impl CfTimeUnits {
43    /// Parse `"<unit> since <reference>"`, for example `"days since 1950-01-01"`.
44    pub fn parse(units: &str) -> Result<Self> {
45        let invalid = || nd_err!(InvalidAttribute, "invalid CF time units '{units}'");
46        let (unit, reference) = units.split_once(" since ").ok_or_else(invalid)?;
47        let step_micros = match unit.trim().to_ascii_lowercase().as_str() {
48            "microseconds" | "microsecond" | "us" => 1,
49            "milliseconds" | "millisecond" | "ms" | "msec" => 1_000,
50            "seconds" | "second" | "secs" | "sec" | "s" => MICROS_PER_SECOND,
51            "minutes" | "minute" | "mins" | "min" => 60 * MICROS_PER_SECOND,
52            "hours" | "hour" | "hrs" | "hr" | "h" => 3_600 * MICROS_PER_SECOND,
53            "days" | "day" | "d" => MICROS_PER_DAY,
54            "weeks" | "week" => 7 * MICROS_PER_DAY,
55            _ => return Err(invalid()),
56        };
57        let reference_micros = parse_datetime(reference.trim()).ok_or_else(invalid)?;
58        Ok(Self {
59            step_micros,
60            reference_micros,
61        })
62    }
63}
64
65/// Days since 1970-01-01 for a proleptic Gregorian date (Howard Hinnant's algorithm).
66pub(crate) fn days_from_civil(year: i64, month: u32, day: u32) -> i64 {
67    let year = if month <= 2 { year - 1 } else { year };
68    let era = year.div_euclid(400);
69    let yoe = year - era * 400;
70    let month = month as i64;
71    let doy = (153 * (if month > 2 { month - 3 } else { month + 9 }) + 2) / 5 + day as i64 - 1;
72    let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
73    era * 146_097 + doe - 719_468
74}
75
76/// The proleptic Gregorian date of a day count since 1970-01-01.
77pub(crate) fn civil_from_days(days: i64) -> (i64, u32, u32) {
78    let z = days + 719_468;
79    let era = z.div_euclid(146_097);
80    let doe = z - era * 146_097;
81    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
82    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
83    let mp = (5 * doy + 2) / 153;
84    let day = (doy - (153 * mp + 2) / 5 + 1) as u32;
85    let month = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
86    let year = yoe + era * 400 + i64::from(month <= 2);
87    (year, month, day)
88}
89
90/// Parse `YYYY-M-D[( |T)h:m[:s[.frac]]][Z|UTC|±hh[:mm]]` to microseconds since the epoch.
91fn parse_datetime(text: &str) -> Option<i64> {
92    let text = text.trim();
93    let (date, rest) = match text.find(['T', ' ']) {
94        Some(i) => (&text[..i], text[i + 1..].trim()),
95        None => (text, ""),
96    };
97    let negative_year = date.starts_with('-');
98    let mut parts = date.trim_start_matches('-').splitn(3, '-');
99    let year: i64 = parts.next()?.parse().ok()?;
100    let year = if negative_year { -year } else { year };
101    let month: u32 = parts.next().map_or(Some(1), |m| m.parse().ok())?;
102    let day: u32 = parts.next().map_or(Some(1), |d| d.parse().ok())?;
103    if !(1..=12).contains(&month) || !(1..=31).contains(&day) {
104        return None;
105    }
106
107    let mut micros = days_from_civil(year, month, day) * MICROS_PER_DAY;
108    let rest = rest.trim_end_matches("UTC").trim();
109    if rest.is_empty() {
110        return Some(micros);
111    }
112    let (time, zone) = match rest.find(['Z', '+', '-']) {
113        Some(i) => (&rest[..i], rest[i..].trim()),
114        None => (rest, ""),
115    };
116    let mut fields = time.trim().split(':');
117    let hour: i64 = fields.next()?.trim().parse().ok()?;
118    let minute: i64 = fields.next().map_or(Some(0), |m| m.parse().ok())?;
119    let seconds: f64 = fields.next().map_or(Some(0.0), |s| s.parse().ok())?;
120    micros += hour * 3_600 * MICROS_PER_SECOND
121        + minute * 60 * MICROS_PER_SECOND
122        + (seconds * MICROS_PER_SECOND as f64).round() as i64;
123
124    if let Some(offset) = zone.strip_prefix(['+', '-']) {
125        let sign = if zone.starts_with('-') { -1 } else { 1 };
126        let digits: String = offset.chars().filter(char::is_ascii_digit).collect();
127        let (h, m) = match digits.len() {
128            1 | 2 => (digits.parse::<i64>().ok()?, 0),
129            4 => (
130                digits[..2].parse::<i64>().ok()?,
131                digits[2..].parse::<i64>().ok()?,
132            ),
133            _ => return None,
134        };
135        micros -= sign * (h * 3_600 + m * 60) * MICROS_PER_SECOND;
136    } else if !(zone.is_empty() || zone == "Z") {
137        return None;
138    }
139    Some(micros)
140}
141
142fn check_calendar(attrs: &Attributes) -> Result<()> {
143    match attrs
144        .get_str(CALENDAR)
145        .map(str::to_ascii_lowercase)
146        .as_deref()
147    {
148        None | Some("standard" | "gregorian" | "proleptic_gregorian") => Ok(()),
149        Some(other) => Err(nd_err!(NotSupported, "CF calendar '{other}'")),
150    }
151}
152
153/// Decode a CF time variable into `Timestamp(Microsecond)`.
154///
155/// Fill values become null first (see [`decode_cf`]). Fractional steps round
156/// to the nearest microsecond. The `units` and `calendar` attributes move to
157/// the returned attributes, for [`encode_cf_time`].
158pub fn decode_cf_time(a: &NdArrowArray) -> Result<(NdArrowArray, Attributes)> {
159    let units = a
160        .attributes()
161        .get_str(UNITS)
162        .ok_or_else(|| nd_err!(InvalidAttribute, "time variable has no 'units' attribute"))?;
163    let parsed = CfTimeUnits::parse(units)?;
164    check_calendar(a.attributes())?;
165    let (masked, _) = decode_cf(a)?;
166
167    let numbers = arrow_cast(masked.values().as_ref(), &DataType::Float64)?;
168    let numbers = numbers.as_primitive::<Float64Type>();
169    let stamps: TimestampMicrosecondArray = numbers
170        .iter()
171        .map(|v| {
172            v.map(|v| parsed.reference_micros + (v * parsed.step_micros as f64).round() as i64)
173        })
174        .collect();
175
176    let mut kept = Attributes::new();
177    let mut moved = Attributes::new();
178    for (key, value) in masked.attributes().iter() {
179        let target = if key == UNITS || key == CALENDAR {
180            &mut moved
181        } else {
182            &mut kept
183        };
184        target.try_insert_array(key, value.clone())?;
185    }
186    Ok((
187        NdArrowArray::try_new(Arc::new(stamps), a.dims().clone())?.with_attributes(kept),
188        moved,
189    ))
190}
191
192/// Encode timestamps as `Float64` numbers of `units`, the inverse of [`decode_cf_time`].
193/// Adds the `units` attribute.
194pub fn encode_cf_time(a: &NdArrowArray, units: &str) -> Result<NdArrowArray> {
195    let parsed = CfTimeUnits::parse(units)?;
196    let stamps = to_micros(a)?;
197    let numbers: Float64Array = stamps
198        .iter()
199        .map(|v| v.map(|v| (v - parsed.reference_micros) as f64 / parsed.step_micros as f64))
200        .collect();
201    let attrs = a.attributes().clone().with(UNITS, units);
202    Ok(NdArrowArray::try_new(Arc::new(numbers), a.dims().clone())?.with_attributes(attrs))
203}
204
205pub(crate) fn to_micros(a: &NdArrowArray) -> Result<TimestampMicrosecondArray> {
206    if !matches!(
207        a.data_type(),
208        DataType::Timestamp(_, _) | DataType::Date32 | DataType::Date64
209    ) {
210        return Err(nd_err!(
211            InvalidArgument,
212            "expected timestamps, got {}",
213            a.data_type()
214        ));
215    }
216    let cast = arrow_cast(a.values().as_ref(), &timestamp_type())?;
217    Ok(cast.as_primitive::<TimestampMicrosecondType>().clone())
218}
219
220/// Calendar periods for [`floor_time`] and [`Frequency`](crate::compute::Frequency).
221#[derive(Debug, Clone, Copy, PartialEq, Eq)]
222pub enum TimePeriod {
223    /// The start of the year.
224    Year,
225    /// The start of the month.
226    Month,
227    /// The start of the day.
228    Day,
229    /// The start of the hour.
230    Hour,
231    /// The start of the minute.
232    Minute,
233    /// The start of the second.
234    Second,
235}
236
237/// Round timestamps down to the start of their `period`, like pandas `floor`.
238/// For regular bins with a reduction, see [`resample`](fn@crate::compute::resample).
239pub fn floor_time(a: &NdArrowArray, period: TimePeriod) -> Result<NdArrowArray> {
240    let stamps = to_micros(a)?;
241    let fixed = |step: i64| move |v: i64| v.div_euclid(step) * step;
242    let out: TimestampMicrosecondArray = match period {
243        TimePeriod::Second => stamps.unary(fixed(MICROS_PER_SECOND)),
244        TimePeriod::Minute => stamps.unary(fixed(60 * MICROS_PER_SECOND)),
245        TimePeriod::Hour => stamps.unary(fixed(3_600 * MICROS_PER_SECOND)),
246        TimePeriod::Day => stamps.unary(fixed(MICROS_PER_DAY)),
247        TimePeriod::Month | TimePeriod::Year => stamps.unary(|v| {
248            let (year, month, _) = civil_from_days(v.div_euclid(MICROS_PER_DAY));
249            let month = if period == TimePeriod::Year { 1 } else { month };
250            days_from_civil(year, month, 1) * MICROS_PER_DAY
251        }),
252    };
253    Ok(NdArrowArray::try_new(Arc::new(out), a.dims().clone())?
254        .with_attributes(a.attributes().clone()))
255}
256
257/// A calendar component of timestamps as `Int32`, for example the month.
258pub fn time_component(a: &NdArrowArray, part: DatePart) -> Result<NdArrowArray> {
259    let stamps: ArrayRef = Arc::new(to_micros(a)?);
260    let out = date_part(stamps.as_ref(), part)?;
261    NdArrowArray::try_new(out, a.dims().clone())
262}
263
264#[cfg(test)]
265mod tests {
266    use arrow::array::{Float32Array, Int32Array};
267    use arrow::datatypes::Int32Type;
268
269    use super::*;
270    use crate::Dimensions;
271    use crate::compute::cf::FILL_VALUE;
272
273    fn micros(year: i64, month: u32, day: u32, hour: i64) -> i64 {
274        days_from_civil(year, month, day) * MICROS_PER_DAY + hour * 3_600 * MICROS_PER_SECOND
275    }
276
277    #[test]
278    fn civil_conversion_round_trips() {
279        assert_eq!(days_from_civil(1970, 1, 1), 0);
280        assert_eq!(days_from_civil(2000, 3, 1), 11_017);
281        for days in [-1_000_000, -719_468, -1, 0, 59, 60, 11_016, 1_000_000] {
282            let (y, m, d) = civil_from_days(days);
283            assert_eq!(days_from_civil(y, m, d), days);
284        }
285    }
286
287    #[test]
288    fn parses_cf_units() {
289        let u = CfTimeUnits::parse("days since 1950-01-01").unwrap();
290        assert_eq!(u.step_micros, MICROS_PER_DAY);
291        assert_eq!(u.reference_micros, micros(1950, 1, 1, 0));
292        let h = CfTimeUnits::parse("hours since 2000-1-1 06:30:00.5 +01:00").unwrap();
293        assert_eq!(
294            h.reference_micros,
295            micros(2000, 1, 1, 5) + 30 * 60 * MICROS_PER_SECOND + 500_000
296        );
297        assert!(CfTimeUnits::parse("days after 1950-01-01").is_err());
298        assert!(CfTimeUnits::parse("fortnights since 1950-01-01").is_err());
299        assert!(CfTimeUnits::parse("seconds since 1970-01-01T00:00:00Z").is_ok());
300    }
301
302    #[test]
303    fn decode_and_encode_round_trip() {
304        let a = NdArrowArray::try_new(
305            Arc::new(Float32Array::from(vec![0.0, 1.5, -999.0])),
306            Dimensions::try_new([("time", 3)]).unwrap(),
307        )
308        .unwrap()
309        .with_attributes(
310            Attributes::new()
311                .with(UNITS, "days since 1950-01-01")
312                .with(CALENDAR, "standard")
313                .with(FILL_VALUE, -999.0f32)
314                .with("long_name", "time"),
315        );
316        let (decoded, encoding) = decode_cf_time(&a).unwrap();
317        assert_eq!(decoded.data_type(), &timestamp_type());
318        let v: Vec<Option<i64>> = decoded
319            .values()
320            .as_primitive::<TimestampMicrosecondType>()
321            .iter()
322            .collect();
323        assert_eq!(
324            v,
325            vec![
326                Some(micros(1950, 1, 1, 0)),
327                Some(micros(1950, 1, 2, 12)),
328                None
329            ]
330        );
331        assert_eq!(decoded.attributes().get_str("long_name"), Some("time"));
332        assert_eq!(encoding.get_str(UNITS), Some("days since 1950-01-01"));
333
334        let encoded = encode_cf_time(&decoded, "days since 1950-01-01").unwrap();
335        let back: Vec<Option<f64>> = encoded
336            .values()
337            .as_primitive::<Float64Type>()
338            .iter()
339            .collect();
340        assert_eq!(back, vec![Some(0.0), Some(1.5), None]);
341    }
342
343    #[test]
344    fn rejects_other_calendars() {
345        let a = NdArrowArray::try_new(
346            Arc::new(Int32Array::from(vec![0])),
347            Dimensions::try_new([("time", 1)]).unwrap(),
348        )
349        .unwrap()
350        .with_attributes(
351            Attributes::new()
352                .with(UNITS, "days since 2000-01-01")
353                .with(CALENDAR, "noleap"),
354        );
355        assert!(decode_cf_time(&a).is_err());
356    }
357
358    #[test]
359    fn floor_and_components() {
360        let stamps = TimestampMicrosecondArray::from(vec![
361            micros(2021, 3, 15, 13),
362            micros(1969, 12, 31, 23),
363        ]);
364        let a = NdArrowArray::try_new(
365            Arc::new(stamps),
366            Dimensions::try_new([("time", 2)]).unwrap(),
367        )
368        .unwrap();
369        let floor = |p| -> Vec<i64> {
370            floor_time(&a, p)
371                .unwrap()
372                .values()
373                .as_primitive::<TimestampMicrosecondType>()
374                .values()
375                .to_vec()
376        };
377        assert_eq!(
378            floor(TimePeriod::Month),
379            vec![micros(2021, 3, 1, 0), micros(1969, 12, 1, 0)]
380        );
381        assert_eq!(
382            floor(TimePeriod::Year),
383            vec![micros(2021, 1, 1, 0), micros(1969, 1, 1, 0)]
384        );
385        assert_eq!(
386            floor(TimePeriod::Day),
387            vec![micros(2021, 3, 15, 0), micros(1969, 12, 31, 0)]
388        );
389        let month = time_component(&a, DatePart::Month).unwrap();
390        assert_eq!(
391            month.values().as_primitive::<Int32Type>().values(),
392            &[3, 12]
393        );
394    }
395}
nd_arrow_array/compute/time.rs

nd_arrow_array/compute/
time.rs
