CoCalc -- test_rolling.py

GitHub Repository: pola-rs/polars
Path: blob/main/py-polars/tests/unit/series/test_rolling.py
⁸⁴²⁴ views
1
from datetime import datetime
2
from typing import Any
3

4
import pytest
5

6
import polars as pl
7
from polars.testing.asserts.series import assert_series_equal
8

9

10
@pytest.fixture
11
def values() -> pl.Series:
12
    return pl.Series([6, 9, 2, 5, 8, 1, 4, 7])
13

14

15
@pytest.fixture
16
def by_col() -> pl.Series:
17
    return pl.Series([0, 0, 1, 1, 2, 2, 3, 3])
18

19

20
@pytest.fixture
21
def by_col_temporal() -> pl.Series:
22
    return pl.Series(
23
        pl.datetime_range(
24
            datetime(2025, 1, 1), datetime(2025, 1, 2), "1h", eager=True
25
        ).head(8)
26
    )
27

28

29
def test_series_rolling_min_by(values: pl.Series, by_col: pl.Series) -> None:
30
    actual = values.rolling_min_by(by_col, "2i")
31
    expected = pl.Series([6, 6, 2, 2, 1, 1, 1, 1])
32
    assert_series_equal(actual, expected)
33

34

35
def test_series_rolling_max_by(values: pl.Series, by_col: pl.Series) -> None:
36
    actual = values.rolling_max_by(by_col, "2i")
37
    expected = pl.Series([9, 9, 9, 9, 8, 8, 8, 8])
38
    assert_series_equal(actual, expected)
39

40

41
def test_series_rolling_sum_by(values: pl.Series, by_col: pl.Series) -> None:
42
    actual = values.rolling_sum_by(by_col, "2i")
43
    expected = pl.Series([15, 15, 22, 22, 16, 16, 20, 20])
44
    assert_series_equal(actual, expected)
45

46

47
def test_series_rolling_mean_by(values: pl.Series, by_col: pl.Series) -> None:
48
    actual = values.rolling_mean_by(by_col, "2i")
49
    expected = pl.Series([7.5, 7.5, 5.5, 5.5, 4.0, 4.0, 5.0, 5.0])
50
    assert_series_equal(actual, expected)
51

52

53
@pytest.mark.parametrize(
54
    ("op", "expected"),
55
    [
56
        ("mean", [None, None, 5.0, 5.0, 4.67, 4.67, 5.0, 5.0]),
57
        ("sum", [0, 0, 5, 5, 14, 14, 20, 20]),
58
        ("min", [None, None, 5, 5, 1, 1, 1, 1]),
59
        ("max", [None, None, 5, 5, 8, 8, 8, 8]),
60
        ("var", [None, None, None, None, 12.33, 12.33, 10.0, 10.0]),
61
        ("std", [None, None, None, None, 3.51, 3.51, 3.16, 3.16]),
62
    ],
63
)
64
def test_series_rolling_by_with_nulls(
65
    values: pl.Series, by_col: pl.Series, op: str, expected: list[float]
66
) -> None:
67
    values[[0, 1, 2]] = None
68
    actual = getattr(values, f"rolling_{op}_by")(by_col, "2i").round(2)
69
    assert_series_equal(actual, pl.Series(expected))
70

71

72
@pytest.mark.parametrize(
73
    ("rank_op", "expected"),
74
    [
75
        ("average", [None, None, None, 1.0, 3.0, 1.0, 2.0, 3.0]),
76
        ("min", [None, None, None, 1, 3, 1, 2, 3]),
77
        ("max", [None, None, None, 1, 3, 1, 2, 3]),
78
        ("dense", [None, None, None, 1, 3, 1, 2, 3]),
79
        ("random", [None, None, None, 1, 3, 1, 2, 3]),
80
    ],
81
)
82
def test_series_rolling_rank_by_with_nulls(
83
    values: pl.Series, by_col: pl.Series, rank_op: Any, expected: list[float]
84
) -> None:
85
    values[[0, 1, 2]] = None
86
    actual = values.rolling_rank_by(by_col, "2i", rank_op).round(2)
87
    dtype = pl.Float64 if rank_op == "average" else pl.UInt32
88
    assert_series_equal(actual, pl.Series(expected, dtype=dtype))
89

90

91
def test_series_rolling_quantile_by_with_nulls(
92
    values: pl.Series, by_col: pl.Series
93
) -> None:
94
    values[[0, 1, 2]] = None
95
    actual = values.rolling_quantile_by(by_col, "2i", quantile=0.35).round(2)
96
    expected = [None, None, 5.0, 5.0, 5.0, 5.0, 4.0, 4.0]
97
    assert_series_equal(actual, pl.Series(expected))
98

99

100
def test_series_rolling_median_by(values: pl.Series, by_col: pl.Series) -> None:
101
    actual = values.rolling_median_by(by_col, "2i")
102
    expected = pl.Series([7.5, 7.5, 5.5, 5.5, 3.5, 3.5, 5.5, 5.5])
103
    assert_series_equal(actual, expected)
104

105

106
def test_series_rolling_std_by(values: pl.Series, by_col: pl.Series) -> None:
107
    actual = values.rolling_std_by(by_col, "2i")
108
    expected = pl.Series([2.12, 2.12, 2.88, 2.88, 3.16, 3.16, 3.16, 3.16])
109
    assert_series_equal(actual, expected, abs_tol=1e-2)
110

111

112
def test_series_rolling_var_by(values: pl.Series, by_col: pl.Series) -> None:
113
    actual = values.rolling_var_by(by_col, "2i")
114
    expected = pl.Series([4.5, 4.5, 8.33, 8.33, 10.0, 10.0, 10.0, 10.0])
115
    assert_series_equal(actual, expected, abs_tol=1e-2)
116

117

118
def test_series_rolling_quantile_by(values: pl.Series, by_col: pl.Series) -> None:
119
    actual = values.rolling_quantile_by(by_col, "2i", quantile=0.5)
120
    expected = pl.Series([9.0, 9.0, 6.0, 6.0, 5.0, 5.0, 7.0, 7.0])
121
    assert_series_equal(actual, expected)
122

123

124
def test_series_rolling_rank_by(values: pl.Series, by_col: pl.Series) -> None:
125
    actual = values.rolling_rank_by(by_col, "2i", method="average")
126
    expected = pl.Series([1.0, 2.0, 1.0, 2.0, 4.0, 1.0, 2.0, 3.0])
127
    assert_series_equal(actual, expected)
128

129

130
def test_series_rolling_min_by_temporal(
131
    values: pl.Series, by_col_temporal: pl.Series
132
) -> None:
133
    actual = values.rolling_min_by(by_col_temporal, "2h")
134
    expected = pl.Series([6, 6, 2, 2, 5, 1, 1, 4])
135
    assert_series_equal(actual, expected)
136

137

138
def test_series_rolling_max_by_temporal(
139
    values: pl.Series, by_col_temporal: pl.Series
140
) -> None:
141
    actual = values.rolling_max_by(by_col_temporal, "2h")
142
    expected = pl.Series([6, 9, 9, 5, 8, 8, 4, 7])
143
    assert_series_equal(actual, expected)
144

145

146
def test_series_rolling_sum_by_temporal(
147
    values: pl.Series, by_col_temporal: pl.Series
148
) -> None:
149
    actual = values.rolling_sum_by(by_col_temporal, "2h")
150
    expected = pl.Series([6, 15, 11, 7, 13, 9, 5, 11])
151
    assert_series_equal(actual, expected)
152

153

154
def test_series_rolling_mean_by_temporal(
155
    values: pl.Series, by_col_temporal: pl.Series
156
) -> None:
157
    actual = values.rolling_mean_by(by_col_temporal, "2h")
158
    expected = pl.Series([6.0, 7.5, 5.5, 3.5, 6.5, 4.5, 2.5, 5.5])
159
    assert_series_equal(actual, expected)
160

161

162
def test_series_rolling_median_by_temporal(
163
    values: pl.Series, by_col_temporal: pl.Series
164
) -> None:
165
    actual = values.rolling_median_by(by_col_temporal, "3h")
166
    expected = pl.Series([6.0, 7.5, 6.0, 5.0, 5.0, 5.0, 4.0, 4.0])
167
    assert_series_equal(actual, expected)
168

169

170
def test_series_rolling_std_by_temporal(
171
    values: pl.Series, by_col_temporal: pl.Series
172
) -> None:
173
    actual = values.rolling_std_by(by_col_temporal, "2h")
174
    expected = pl.Series([None, 2.12, 4.94, 2.12, 2.12, 4.94, 2.12, 2.12])
175
    assert_series_equal(actual, expected, abs_tol=1e-2)
176

177

178
def test_series_rolling_var_by_temporal(
179
    values: pl.Series, by_col_temporal: pl.Series
180
) -> None:
181
    actual = values.rolling_var_by(by_col_temporal, "2h")
182
    expected = pl.Series([None, 4.5, 24.5, 4.5, 4.5, 24.5, 4.5, 4.5])
183
    assert_series_equal(actual, expected, abs_tol=1e-2)
184

185

186
def test_series_rolling_quantile_by_temporal(
187
    values: pl.Series, by_col_temporal: pl.Series
188
) -> None:
189
    actual = values.rolling_quantile_by(by_col_temporal, "2h", quantile=0.5)
190
    expected = pl.Series([6.0, 9.0, 9.0, 5.0, 8.0, 8.0, 4.0, 7.0])
191
    assert_series_equal(actual, expected)
192

193

194
def test_series_rolling_rank_by_temporal(
195
    values: pl.Series, by_col_temporal: pl.Series
196
) -> None:
197
    actual = values.rolling_rank_by(by_col_temporal, "2h", method="average")
198
    expected = pl.Series([1.0, 2.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0])
199
    assert_series_equal(actual, expected)
200

201
Product

Resources

Company
