polars_plan/dsl/function_expr/
rolling_by.rs

1use polars_time::chunkedarray::*;
2
3use super::*;
4
5#[derive(Clone, PartialEq, Debug)]
6#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7pub enum RollingFunctionBy {
8    MinBy(RollingOptionsDynamicWindow),
9    MaxBy(RollingOptionsDynamicWindow),
10    MeanBy(RollingOptionsDynamicWindow),
11    SumBy(RollingOptionsDynamicWindow),
12    QuantileBy(RollingOptionsDynamicWindow),
13    VarBy(RollingOptionsDynamicWindow),
14    StdBy(RollingOptionsDynamicWindow),
15}
16
17impl Display for RollingFunctionBy {
18    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
19        use RollingFunctionBy::*;
20
21        let name = match self {
22            MinBy(_) => "rolling_min_by",
23            MaxBy(_) => "rolling_max_by",
24            MeanBy(_) => "rolling_mean_by",
25            SumBy(_) => "rolling_sum_by",
26            QuantileBy(_) => "rolling_quantile_by",
27            VarBy(_) => "rolling_var_by",
28            StdBy(_) => "rolling_std_by",
29        };
30
31        write!(f, "{name}")
32    }
33}
34
35impl Hash for RollingFunctionBy {
36    fn hash<H: Hasher>(&self, state: &mut H) {
37        std::mem::discriminant(self).hash(state);
38    }
39}
40
41pub(super) fn rolling_min_by(
42    s: &[Column],
43    options: RollingOptionsDynamicWindow,
44) -> PolarsResult<Column> {
45    // @scalar-opt
46    s[0].as_materialized_series()
47        .rolling_min_by(s[1].as_materialized_series(), options)
48        .map(Column::from)
49}
50
51pub(super) fn rolling_max_by(
52    s: &[Column],
53    options: RollingOptionsDynamicWindow,
54) -> PolarsResult<Column> {
55    // @scalar-opt
56    s[0].as_materialized_series()
57        .rolling_max_by(s[1].as_materialized_series(), options)
58        .map(Column::from)
59}
60
61pub(super) fn rolling_mean_by(
62    s: &[Column],
63    options: RollingOptionsDynamicWindow,
64) -> PolarsResult<Column> {
65    // @scalar-opt
66    s[0].as_materialized_series()
67        .rolling_mean_by(s[1].as_materialized_series(), options)
68        .map(Column::from)
69}
70
71pub(super) fn rolling_sum_by(
72    s: &[Column],
73    options: RollingOptionsDynamicWindow,
74) -> PolarsResult<Column> {
75    // @scalar-opt
76    s[0].as_materialized_series()
77        .rolling_sum_by(s[1].as_materialized_series(), options)
78        .map(Column::from)
79}
80
81pub(super) fn rolling_quantile_by(
82    s: &[Column],
83    options: RollingOptionsDynamicWindow,
84) -> PolarsResult<Column> {
85    // @scalar-opt
86    s[0].as_materialized_series()
87        .rolling_quantile_by(s[1].as_materialized_series(), options)
88        .map(Column::from)
89}
90
91pub(super) fn rolling_var_by(
92    s: &[Column],
93    options: RollingOptionsDynamicWindow,
94) -> PolarsResult<Column> {
95    // @scalar-opt
96    s[0].as_materialized_series()
97        .rolling_var_by(s[1].as_materialized_series(), options)
98        .map(Column::from)
99}
100
101pub(super) fn rolling_std_by(
102    s: &[Column],
103    options: RollingOptionsDynamicWindow,
104) -> PolarsResult<Column> {
105    // @scalar-opt
106    s[0].as_materialized_series()
107        .rolling_std_by(s[1].as_materialized_series(), options)
108        .map(Column::from)
109}