Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
 Path: blob/main/py-polars/tests/unit/series/test_rolling.py
6939 views
1
from datetime import datetime

2


3
import pytest

4


5
import polars as pl

6
from polars.testing.asserts.series import assert_series_equal

7


8


9
@pytest.fixture

10
def values() -> pl.Series:

11
    return pl.Series([6, 9, 2, 5, 8, 1, 4, 7])

12


13


14
@pytest.fixture

15
def by_col() -> pl.Series:

16
    return pl.Series([0, 0, 1, 1, 2, 2, 3, 3])

17


18


19
@pytest.fixture

20
def by_col_temporal() -> pl.Series:

21
    return pl.Series(

22
        pl.datetime_range(

23
            datetime(2025, 1, 1), datetime(2025, 1, 2), "1h", eager=True

24
        ).head(8)

25
    )

26


27


28
def test_series_rolling_min_by(values: pl.Series, by_col: pl.Series) -> None:

29
    actual = values.rolling_min_by(by_col, "2i")

30
    expected = pl.Series([6, 6, 2, 2, 1, 1, 1, 1])

31
    assert_series_equal(actual, expected)

32


33


34
def test_series_rolling_max_by(values: pl.Series, by_col: pl.Series) -> None:

35
    actual = values.rolling_max_by(by_col, "2i")

36
    expected = pl.Series([9, 9, 9, 9, 8, 8, 8, 8])

37
    assert_series_equal(actual, expected)

38


39


40
def test_series_rolling_sum_by(values: pl.Series, by_col: pl.Series) -> None:

41
    actual = values.rolling_sum_by(by_col, "2i")

42
    expected = pl.Series([15, 15, 22, 22, 16, 16, 20, 20])

43
    assert_series_equal(actual, expected)

44


45


46
def test_series_rolling_mean_by(values: pl.Series, by_col: pl.Series) -> None:

47
    actual = values.rolling_mean_by(by_col, "2i")

48
    expected = pl.Series([7.5, 7.5, 5.5, 5.5, 4.0, 4.0, 5.0, 5.0])

49
    assert_series_equal(actual, expected)

50


51


52
def test_series_rolling_median_by(values: pl.Series, by_col: pl.Series) -> None:

53
    actual = values.rolling_median_by(by_col, "2i")

54
    expected = pl.Series([7.5, 7.5, 5.5, 5.5, 3.5, 3.5, 5.5, 5.5])

55
    assert_series_equal(actual, expected)

56


57


58
def test_series_rolling_std_by(values: pl.Series, by_col: pl.Series) -> None:

59
    actual = values.rolling_std_by(by_col, "2i")

60
    expected = pl.Series([2.12, 2.12, 2.88, 2.88, 3.16, 3.16, 3.16, 3.16])

61
    assert_series_equal(actual, expected, abs_tol=1e-2)

62


63


64
def test_series_rolling_var_by(values: pl.Series, by_col: pl.Series) -> None:

65
    actual = values.rolling_var_by(by_col, "2i")

66
    expected = pl.Series([4.5, 4.5, 8.33, 8.33, 10.0, 10.0, 10.0, 10.0])

67
    assert_series_equal(actual, expected, abs_tol=1e-2)

68


69


70
def test_series_rolling_quantile_by(values: pl.Series, by_col: pl.Series) -> None:

71
    actual = values.rolling_quantile_by(by_col, "2i", quantile=0.5)

72
    expected = pl.Series([9.0, 9.0, 6.0, 6.0, 5.0, 5.0, 7.0, 7.0])

73
    assert_series_equal(actual, expected)

74


75


76
def test_series_rolling_min_by_temporal(

77
    values: pl.Series, by_col_temporal: pl.Series

78
) -> None:

79
    actual = values.rolling_min_by(by_col_temporal, "2h")

80
    expected = pl.Series([6, 6, 2, 2, 5, 1, 1, 4])

81
    assert_series_equal(actual, expected)

82


83


84
def test_series_rolling_max_by_temporal(

85
    values: pl.Series, by_col_temporal: pl.Series

86
) -> None:

87
    actual = values.rolling_max_by(by_col_temporal, "2h")

88
    expected = pl.Series([6, 9, 9, 5, 8, 8, 4, 7])

89
    assert_series_equal(actual, expected)

90


91


92
def test_series_rolling_sum_by_temporal(

93
    values: pl.Series, by_col_temporal: pl.Series

94
) -> None:

95
    actual = values.rolling_sum_by(by_col_temporal, "2h")

96
    expected = pl.Series([6, 15, 11, 7, 13, 9, 5, 11])

97
    assert_series_equal(actual, expected)

98


99


100
def test_series_rolling_mean_by_temporal(

101
    values: pl.Series, by_col_temporal: pl.Series

102
) -> None:

103
    actual = values.rolling_mean_by(by_col_temporal, "2h")

104
    expected = pl.Series([6.0, 7.5, 5.5, 3.5, 6.5, 4.5, 2.5, 5.5])

105
    assert_series_equal(actual, expected)

106


107


108
def test_series_rolling_median_by_temporal(

109
    values: pl.Series, by_col_temporal: pl.Series

110
) -> None:

111
    actual = values.rolling_median_by(by_col_temporal, "3h")

112
    expected = pl.Series([6.0, 7.5, 6.0, 5.0, 5.0, 5.0, 4.0, 4.0])

113
    assert_series_equal(actual, expected)

114


115


116
def test_series_rolling_std_by_temporal(

117
    values: pl.Series, by_col_temporal: pl.Series

118
) -> None:

119
    actual = values.rolling_std_by(by_col_temporal, "2h")

120
    expected = pl.Series([None, 2.12, 4.94, 2.12, 2.12, 4.94, 2.12, 2.12])

121
    assert_series_equal(actual, expected, abs_tol=1e-2)

122


123


124
def test_series_rolling_var_by_temporal(

125
    values: pl.Series, by_col_temporal: pl.Series

126
) -> None:

127
    actual = values.rolling_var_by(by_col_temporal, "2h")

128
    expected = pl.Series([None, 4.5, 24.5, 4.5, 4.5, 24.5, 4.5, 4.5])

129
    assert_series_equal(actual, expected, abs_tol=1e-2)

130


131


132
def test_series_rolling_quantile_by_temporal(

133
    values: pl.Series, by_col_temporal: pl.Series

134
) -> None:

135
    actual = values.rolling_quantile_by(by_col_temporal, "2h", quantile=0.5)

136
    expected = pl.Series([6.0, 9.0, 9.0, 5.0, 8.0, 8.0, 4.0, 7.0])

137
    assert_series_equal(actual, expected)

138


139