CoCalc -- test_date

GitHub Repository: pola-rs/polars
Path: blob/main/py-polars/tests/unit/functions/range/test_date_ranges.py
⁷⁸⁸⁴ views
1
from __future__ import annotations
2

3
from datetime import date
4
from typing import TYPE_CHECKING
5

6
import pandas as pd
7
import pytest
8

9
import polars as pl
10
from polars.exceptions import ComputeError
11
from polars.testing import assert_frame_equal, assert_series_equal
12

13
if TYPE_CHECKING:
14
    from polars._typing import ClosedInterval
15

16

17
def test_date_ranges_lazy_with_literals() -> None:
18
    df = pl.DataFrame({"misc": ["x"]}).with_columns(
19
        pl.date_ranges(
20
            start=date(2000, 1, 1),
21
            end=date(2023, 8, 31),
22
            interval="987d",
23
            eager=False,
24
        ).alias("dts")
25
    )
26
    assert df.rows() == [
27
        (
28
            "x",
29
            [
30
                date(2000, 1, 1),
31
                date(2002, 9, 14),
32
                date(2005, 5, 28),
33
                date(2008, 2, 9),
34
                date(2010, 10, 23),
35
                date(2013, 7, 6),
36
                date(2016, 3, 19),
37
                date(2018, 12, 1),
38
                date(2021, 8, 14),
39
            ],
40
        )
41
    ]
42
    assert (
43
        df.rows()[0][1]
44
        == pd.date_range(
45
            date(2000, 1, 1), date(2023, 12, 31), freq="987d"
46
        ).date.tolist()
47
    )
48

49

50
@pytest.mark.parametrize("low", ["start", pl.col("start")])
51
@pytest.mark.parametrize("high", ["stop", pl.col("stop")])
52
def test_date_ranges_lazy_with_expressions(
53
    low: str | pl.Expr, high: str | pl.Expr
54
) -> None:
55
    lf = pl.LazyFrame(
56
        {
57
            "start": [date(2015, 6, 30)],
58
            "stop": [date(2022, 12, 31)],
59
        }
60
    )
61

62
    result = lf.with_columns(
63
        pl.date_ranges(start=low, end=high, interval="678d", eager=False).alias("dts")
64
    )
65

66
    assert result.collect().rows() == [
67
        (
68
            date(2015, 6, 30),
69
            date(2022, 12, 31),
70
            [
71
                date(2015, 6, 30),
72
                date(2017, 5, 8),
73
                date(2019, 3, 17),
74
                date(2021, 1, 23),
75
                date(2022, 12, 2),
76
            ],
77
        )
78
    ]
79

80
    df = pl.DataFrame(
81
        {
82
            "start": [date(2000, 1, 1), date(2022, 6, 1)],
83
            "stop": [date(2000, 1, 2), date(2022, 6, 2)],
84
        }
85
    )
86

87
    result_df = df.with_columns(
88
        pl.date_ranges(start=low, end=high, interval="1d").alias("dts")
89
    )
90

91
    assert result_df.to_dict(as_series=False) == {
92
        "start": [date(2000, 1, 1), date(2022, 6, 1)],
93
        "stop": [date(2000, 1, 2), date(2022, 6, 2)],
94
        "dts": [
95
            [date(2000, 1, 1), date(2000, 1, 2)],
96
            [date(2022, 6, 1), date(2022, 6, 2)],
97
        ],
98
    }
99

100

101
def test_date_ranges_single_row_lazy_7110() -> None:
102
    df = pl.DataFrame(
103
        {
104
            "name": ["A"],
105
            "from": [date(2020, 1, 1)],
106
            "to": [date(2020, 1, 2)],
107
        }
108
    )
109
    result = df.with_columns(
110
        pl.date_ranges(
111
            start=pl.col("from"),
112
            end=pl.col("to"),
113
            interval="1d",
114
            eager=False,
115
        ).alias("date_range")
116
    )
117
    expected = pl.DataFrame(
118
        {
119
            "name": ["A"],
120
            "from": [date(2020, 1, 1)],
121
            "to": [date(2020, 1, 2)],
122
            "date_range": [[date(2020, 1, 1), date(2020, 1, 2)]],
123
        }
124
    )
125
    assert_frame_equal(result, expected)
126

127

128
def test_date_ranges_eager() -> None:
129
    start = pl.Series("start", [date(2022, 1, 1), date(2022, 1, 2)])
130
    end = pl.Series("end", [date(2022, 1, 4), date(2022, 1, 3)])
131

132
    result = pl.date_ranges(start=start, end=end, eager=True)
133

134
    expected = pl.Series(
135
        "start",
136
        [
137
            [date(2022, 1, 1), date(2022, 1, 2), date(2022, 1, 3), date(2022, 1, 4)],
138
            [date(2022, 1, 2), date(2022, 1, 3)],
139
        ],
140
    )
141
    assert_series_equal(result, expected)
142

143

144
def test_date_ranges_broadcasting() -> None:
145
    df = pl.DataFrame({"dates": [date(2021, 1, 1), date(2021, 1, 2), date(2021, 1, 3)]})
146
    result = df.select(
147
        pl.date_ranges(start="dates", end=date(2021, 1, 3)).alias("end"),
148
        pl.date_ranges(start=date(2021, 1, 1), end="dates").alias("start"),
149
    )
150
    expected = pl.DataFrame(
151
        {
152
            "end": [
153
                [date(2021, 1, 1), date(2021, 1, 2), date(2021, 1, 3)],
154
                [date(2021, 1, 2), date(2021, 1, 3)],
155
                [date(2021, 1, 3)],
156
            ],
157
            "start": [
158
                [date(2021, 1, 1)],
159
                [date(2021, 1, 1), date(2021, 1, 2)],
160
                [date(2021, 1, 1), date(2021, 1, 2), date(2021, 1, 3)],
161
            ],
162
        }
163
    )
164
    assert_frame_equal(result, expected)
165

166

167
def test_date_ranges_broadcasting_fail() -> None:
168
    start = pl.Series([date(2021, 1, 1), date(2021, 1, 2), date(2021, 1, 3)])
169
    end = pl.Series([date(2021, 1, 2), date(2021, 1, 3)])
170

171
    with pytest.raises(
172
        ComputeError, match=r"lengths of `start` \(3\) and `end` \(2\) do not match"
173
    ):
174
        pl.date_ranges(start=start, end=end, eager=True)
175

176

177
@pytest.mark.parametrize(
178
    ("closed", "expected"),
179
    [
180
        (
181
            "both",
182
            [
183
                [
184
                    date(2025, 1, 1),
185
                    date(2025, 1, 3),
186
                    date(2025, 1, 5),
187
                    date(2025, 1, 7),
188
                ],
189
                [
190
                    date(2025, 1, 11),
191
                    date(2025, 1, 13),
192
                    date(2025, 1, 15),
193
                    date(2025, 1, 17),
194
                ],
195
            ],
196
        ),
197
        (
198
            "left",
199
            [
200
                [
201
                    date(2025, 1, 1),
202
                    date(2025, 1, 3),
203
                    date(2025, 1, 5),
204
                    date(2025, 1, 7),
205
                ],
206
                [
207
                    date(2025, 1, 11),
208
                    date(2025, 1, 13),
209
                    date(2025, 1, 15),
210
                    date(2025, 1, 17),
211
                ],
212
            ],
213
        ),
214
        (
215
            "right",
216
            [
217
                [date(2025, 1, 3), date(2025, 1, 5), date(2025, 1, 7)],
218
                [date(2025, 1, 13), date(2025, 1, 15), date(2025, 1, 17)],
219
            ],
220
        ),
221
        (
222
            "none",
223
            [
224
                [date(2025, 1, 3), date(2025, 1, 5), date(2025, 1, 7)],
225
                [date(2025, 1, 13), date(2025, 1, 15), date(2025, 1, 17)],
226
            ],
227
        ),
228
    ],
229
)
230
def test_date_ranges_start_end_interval_forwards(
231
    closed: ClosedInterval,
232
    expected: list[date],
233
) -> None:
234
    df = pl.DataFrame(
235
        {
236
            "start": [date(2025, 1, 1), date(2025, 1, 11)],
237
            "end": [date(2025, 1, 8), date(2025, 1, 18)],
238
        }
239
    )
240
    result = df.select(
241
        dates=pl.date_ranges(start="start", end="end", interval="2d", closed=closed)
242
    )
243
    assert_frame_equal(result, pl.Series("dates", expected).to_frame())
244

245
Product

Resources

Company
