polars_time/windows/
index_space.rs

1//! The `i64` space a temporal group-by computes its windows in.
2//!
3//! Both engines cast the index column to a `Datetime` before computing windows, and work in
4//! that dtype's physical `i64` values: `Datetime` in its own time unit and zone, `Date` as
5//! microseconds, integers reinterpreted as nanoseconds. [`IndexSpace`] holds that mapping in
6//! one place for the in-memory engine and the streaming nodes.
7
8use polars_arrow::legacy::time_zone::Tz;
9use polars_core::prelude::*;
10
11/// The `i64` space a temporal group-by computes its windows in: the unit and zone of the
12/// `Datetime` the index is cast to.
13#[derive(Clone)]
14pub struct IndexSpace {
15    pub time_unit: TimeUnit,
16    time_zone: Option<TimeZone>,
17    tz: Option<Tz>,
18    index_dtype: DataType,
19}
20
21impl std::fmt::Debug for IndexSpace {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        f.debug_struct("IndexSpace")
24            .field("time_unit", &self.time_unit)
25            .field("time_zone", &self.time_zone)
26            .field("index_dtype", &self.index_dtype)
27            .finish()
28    }
29}
30
31impl IndexSpace {
32    /// The space for a `group_by_dynamic` index of `index_dtype`, which must be a `Datetime`,
33    /// `Date` or signed integer.
34    ///
35    /// A zone `chrono-tz` cannot parse, which needs `POLARS_IGNORE_TIMEZONE_PARSE_ERROR=1` or
36    /// the `timezones` feature to be off, is kept in `time_zone` but ignored for the
37    /// arithmetic: the engines compute such windows as if the index had no zone, and so does
38    /// this.
39    pub fn dynamic(index_dtype: &DataType) -> PolarsResult<Self> {
40        Self::new(index_dtype, false)
41    }
42
43    /// The space for a `rolling` index of `index_dtype`, which may also be an unsigned integer.
44    /// See [`Self::dynamic`] for the zone handling.
45    pub fn rolling(index_dtype: &DataType) -> PolarsResult<Self> {
46        Self::new(index_dtype, true)
47    }
48
49    fn new(index_dtype: &DataType, allow_unsigned: bool) -> PolarsResult<Self> {
50        let DataType::Datetime(time_unit, time_zone) =
51            window_datetime_dtype(index_dtype, allow_unsigned)?
52        else {
53            unreachable!()
54        };
55        #[cfg(feature = "timezones")]
56        let tz = time_zone.as_ref().and_then(|tz| tz.parse::<Tz>().ok());
57        #[cfg(not(feature = "timezones"))]
58        let tz = None;
59        Ok(Self {
60            time_unit,
61            time_zone,
62            tz,
63            index_dtype: index_dtype.clone(),
64        })
65    }
66
67    /// The zone the arithmetic runs in, as `Duration` and `Window` take it. `None` when the
68    /// index has no zone or one the engines ignore, see [`Self::dynamic`].
69    pub fn tz(&self) -> Option<&Tz> {
70        self.tz.as_ref()
71    }
72
73    /// The zone of the `Datetime` dtype of this space.
74    pub fn time_zone(&self) -> Option<&TimeZone> {
75        self.time_zone.as_ref()
76    }
77
78    /// The `Datetime` dtype of this space.
79    pub fn window_dtype(&self) -> DataType {
80        DataType::Datetime(self.time_unit, self.time_zone.clone())
81    }
82
83    /// `index`, a column of the dtype this space was made for, cast to [`Self::window_dtype`].
84    pub fn cast_to_space(&self, index: &Column) -> PolarsResult<Column> {
85        debug_assert_eq!(index.dtype(), &self.index_dtype);
86        match &self.index_dtype {
87            DataType::Datetime(_, _) => Ok(index.clone()),
88            DataType::Int32 | DataType::UInt32 | DataType::UInt64 => {
89                index.cast(&DataType::Int64)?.cast(&self.window_dtype())
90            },
91            _ => index.cast(&self.window_dtype()),
92        }
93    }
94
95    /// `column`, in [`Self::window_dtype`], cast back to the dtype this space was made for.
96    pub fn cast_from_space(&self, column: &Column) -> PolarsResult<Column> {
97        debug_assert_eq!(column.dtype(), &self.window_dtype());
98        match &self.index_dtype {
99            DataType::Datetime(_, _) => Ok(column.clone()),
100            dt if dt.is_integer() => column.cast(&DataType::Int64)?.cast(dt),
101            dt => column.cast(dt),
102        }
103    }
104
105    /// `column`, in [`Self::window_dtype`], cast to the dtype of a `group_by_dynamic` window
106    /// boundary: the index dtype, except that boundaries of a `Date` index stay a `Datetime`
107    /// because a window need not start on a day.
108    pub fn cast_to_boundary(&self, column: &Column) -> PolarsResult<Column> {
109        match &self.index_dtype {
110            DataType::Date => Ok(column.clone()),
111            _ => self.cast_from_space(column),
112        }
113    }
114}
115
116/// The `Datetime` dtype an index column is cast to before its windows are computed: `Date`
117/// becomes microseconds since the epoch and integers are reinterpreted as nanoseconds.
118fn window_datetime_dtype(index_dtype: &DataType, allow_unsigned: bool) -> PolarsResult<DataType> {
119    use DataType::*;
120    Ok(match index_dtype {
121        Datetime(_, _) => index_dtype.clone(),
122        Date => Datetime(TimeUnit::Microseconds, None),
123        Int32 | Int64 => Datetime(TimeUnit::Nanoseconds, None),
124        UInt32 | UInt64 if allow_unsigned => Datetime(TimeUnit::Nanoseconds, None),
125        dt if allow_unsigned => polars_bail!(
126            ComputeError:
127            "expected any of the following dtypes: {{ Date, Datetime, Int32, Int64, UInt32, UInt64 }}, got {}",
128            dt
129        ),
130        dt => polars_bail!(
131            ComputeError:
132            "expected any of the following dtypes: {{ Date, Datetime, Int32, Int64 }}, got {}",
133            dt
134        ),
135    })
136}
polars_time/windows/index_space.rs

polars_time/windows/
index_space.rs
