CoCalc -- test_time

GitHub Repository: pola-rs/polars
Path: blob/main/py-polars/tests/unit/functions/range/test_time_range.py
⁶⁹³⁹ views
1
from __future__ import annotations
2

3
from datetime import time, timedelta
4
from typing import TYPE_CHECKING
5

6
import pytest
7

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

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

15

16
def test_time_range_schema() -> None:
17
    df = pl.DataFrame({"start": [time(1)], "end": [time(1, 30)]}).lazy()
18
    result = df.with_columns(time_range=pl.time_ranges(pl.col("start"), pl.col("end")))
19
    expected_schema = {"start": pl.Time, "end": pl.Time, "time_range": pl.List(pl.Time)}
20
    assert result.collect_schema() == expected_schema
21
    assert result.collect().schema == expected_schema
22

23

24
def test_time_ranges_eager() -> None:
25
    start = pl.Series("start", [time(9, 0), time(10, 0)])
26
    end = pl.Series("end", [time(12, 0), time(11, 0)])
27

28
    result = pl.time_ranges(start, end, eager=True)
29

30
    expected = pl.Series(
31
        "start",
32
        [
33
            [time(9, 0), time(10, 0), time(11, 0), time(12, 0)],
34
            [time(10, 0), time(11, 0)],
35
        ],
36
    )
37
    assert_series_equal(result, expected)
38

39

40
def test_time_range_eager_explode() -> None:
41
    result = pl.time_range(time(9, 0), time(11, 0), eager=True)
42
    expected = pl.Series("literal", [time(9, 0), time(10, 0), time(11, 0)])
43
    assert_series_equal(result, expected)
44

45

46
def test_time_range_input_shape_empty() -> None:
47
    empty = pl.Series(dtype=pl.Time)
48
    single = pl.Series([time(12, 0)])
49

50
    with pytest.raises(ShapeError):
51
        pl.time_range(empty, single, eager=True)
52
    with pytest.raises(ShapeError):
53
        pl.time_range(single, empty, eager=True)
54
    with pytest.raises(ShapeError):
55
        pl.time_range(empty, empty, eager=True)
56

57

58
def test_time_range_input_shape_multiple_values() -> None:
59
    single = pl.Series([time(12, 0)])
60
    multiple = pl.Series([time(11, 0), time(12, 0)])
61

62
    with pytest.raises(ShapeError):
63
        pl.time_range(multiple, single, eager=True)
64
    with pytest.raises(ShapeError):
65
        pl.time_range(single, multiple, eager=True)
66
    with pytest.raises(ShapeError):
67
        pl.time_range(multiple, multiple, eager=True)
68

69

70
def test_time_range_start_equals_end() -> None:
71
    t = time(12, 0)
72

73
    result = pl.time_range(t, t, closed="both", eager=True)
74

75
    expected = pl.Series("literal", [t])
76
    assert_series_equal(result, expected)
77

78

79
@pytest.mark.parametrize("closed", ["left", "right", "none"])
80
def test_time_range_start_equals_end_open(closed: ClosedInterval) -> None:
81
    t = time(12, 0)
82

83
    result = pl.time_range(t, t, closed=closed, eager=True)
84

85
    expected = pl.Series("literal", dtype=pl.Time)
86
    assert_series_equal(result, expected)
87

88

89
def test_time_range_start_later_than_end() -> None:
90
    result = pl.time_range(time(12), time(11), eager=True)
91
    expected = pl.Series("literal", dtype=pl.Time)
92
    assert_series_equal(result, expected)
93

94

95
@pytest.mark.parametrize("interval", [timedelta(0), timedelta(minutes=-10)])
96
def test_time_range_invalid_step(interval: timedelta) -> None:
97
    with pytest.raises(ComputeError, match="`interval` must be positive"):
98
        pl.time_range(time(11), time(12), interval=interval, eager=True)
99

100

101
def test_time_range_lit_lazy() -> None:
102
    tm = pl.select(
103
        pl.time_range(
104
            start=time(1, 2, 3),
105
            end=time(23, 59, 59),
106
            interval="5h45m10s333ms",
107
            closed="right",
108
        ).alias("tm")
109
    )
110

111
    assert tm["tm"].to_list() == [
112
        time(6, 47, 13, 333000),
113
        time(12, 32, 23, 666000),
114
        time(18, 17, 33, 999000),
115
    ]
116

117
    # validate unset start/end
118
    tm = pl.select(pl.time_range(interval="5h45m10s333ms").alias("tm"))
119
    assert tm["tm"].to_list() == [
120
        time(0, 0),
121
        time(5, 45, 10, 333000),
122
        time(11, 30, 20, 666000),
123
        time(17, 15, 30, 999000),
124
        time(23, 0, 41, 332000),
125
    ]
126

127
    tm = pl.select(
128
        pl.time_range(start=pl.lit(time(23, 59, 59, 999980)), interval="10000ns").alias(
129
            "tm"
130
        )
131
    )
132
    assert tm["tm"].to_list() == [
133
        time(23, 59, 59, 999980),
134
        time(23, 59, 59, 999990),
135
    ]
136

137

138
def test_time_range_lit_eager() -> None:
139
    eager = True
140
    tm = pl.select(
141
        pl.time_range(
142
            start=time(1, 2, 3),
143
            end=time(23, 59, 59),
144
            interval="5h45m10s333ms",
145
            closed="right",
146
            eager=eager,
147
        ).alias("tm")
148
    )
149
    if not eager:
150
        tm = tm.select(pl.col("tm").explode())
151
    assert tm["tm"].to_list() == [
152
        time(6, 47, 13, 333000),
153
        time(12, 32, 23, 666000),
154
        time(18, 17, 33, 999000),
155
    ]
156

157
    # validate unset start/end
158
    tm = pl.select(
159
        pl.time_range(
160
            interval="5h45m10s333ms",
161
            eager=eager,
162
        ).alias("tm")
163
    )
164
    if not eager:
165
        tm = tm.select(pl.col("tm").explode())
166
    assert tm["tm"].to_list() == [
167
        time(0, 0),
168
        time(5, 45, 10, 333000),
169
        time(11, 30, 20, 666000),
170
        time(17, 15, 30, 999000),
171
        time(23, 0, 41, 332000),
172
    ]
173

174
    tm = pl.select(
175
        pl.time_range(
176
            start=pl.lit(time(23, 59, 59, 999980)),
177
            interval="10000ns",
178
            eager=eager,
179
        ).alias("tm")
180
    )
181
    tm = tm.select(pl.col("tm").explode())
182
    assert tm["tm"].to_list() == [
183
        time(23, 59, 59, 999980),
184
        time(23, 59, 59, 999990),
185
    ]
186

187

188
def test_time_range_expr() -> None:
189
    df = pl.DataFrame(
190
        {
191
            "start": pl.time_range(interval="6h", eager=True),
192
            "stop": pl.time_range(start=time(2, 59), interval="5h59m", eager=True),
193
        }
194
    ).with_columns(intervals=pl.time_ranges("start", pl.col("stop"), interval="1h29m"))
195
    # shape: (4, 3)
196
    # ┌──────────┬──────────┬────────────────────────────────┐
197
    # │ start    ┆ stop     ┆ intervals                      │
198
    # │ ---      ┆ ---      ┆ ---                            │
199
    # │ time     ┆ time     ┆ list[time]                     │
200
    # ╞══════════╪══════════╪════════════════════════════════╡
201
    # │ 00:00:00 ┆ 02:59:00 ┆ [00:00:00, 01:29:00, 02:58:00] │
202
    # │ 06:00:00 ┆ 08:58:00 ┆ [06:00:00, 07:29:00, 08:58:00] │
203
    # │ 12:00:00 ┆ 14:57:00 ┆ [12:00:00, 13:29:00]           │
204
    # │ 18:00:00 ┆ 20:56:00 ┆ [18:00:00, 19:29:00]           │
205
    # └──────────┴──────────┴────────────────────────────────┘
206
    assert df.rows() == [
207
        (time(0, 0), time(2, 59), [time(0, 0), time(1, 29), time(2, 58)]),
208
        (time(6, 0), time(8, 58), [time(6, 0), time(7, 29), time(8, 58)]),
209
        (time(12, 0), time(14, 57), [time(12, 0), time(13, 29)]),
210
        (time(18, 0), time(20, 56), [time(18, 0), time(19, 29)]),
211
    ]
212

213

214
def test_time_range_name() -> None:
215
    expected_name = "literal"
216
    result_eager = pl.time_range(time(10), time(12), eager=True)
217
    assert result_eager.name == expected_name
218

219
    expected_name = "s1"
220
    result_lazy = pl.select(
221
        pl.time_range(
222
            pl.lit(time(10)).alias("s1"), pl.lit(time(12)).alias("s2"), eager=False
223
        )
224
    ).to_series()
225
    assert result_lazy.name == expected_name
226

227

228
def test_time_ranges_broadcasting() -> None:
229
    df = pl.DataFrame({"time": [time(10, 0), time(11, 0), time(12, 0)]})
230
    result = df.select(
231
        pl.time_ranges(start="time", end=time(12, 0)).alias("end"),
232
        pl.time_ranges(start=time(10, 0), end="time").alias("start"),
233
    )
234
    expected = pl.DataFrame(
235
        {
236
            "end": [
237
                [time(10, 0), time(11, 0), time(12, 0)],
238
                [time(11, 0), time(12, 0)],
239
                [time(12, 0)],
240
            ],
241
            "start": [
242
                [time(10, 0)],
243
                [time(10, 0), time(11, 0)],
244
                [time(10, 0), time(11, 0), time(12, 0)],
245
            ],
246
        }
247
    )
248
    assert_frame_equal(result, expected)
249

250

251
def test_time_ranges_mismatched_chunks() -> None:
252
    s1 = pl.Series("s1", [time(10), time(11)])
253
    s1.append(pl.Series([time(12)]))
254

255
    s2 = pl.Series("s2", [time(12)])
256
    s2.append(pl.Series([time(12), time(12)]))
257

258
    result = pl.time_ranges(s1, s2, eager=True)
259
    expected = pl.Series(
260
        "s1",
261
        [
262
            [time(10, 0), time(11, 0), time(12, 0)],
263
            [time(11, 0), time(12, 0)],
264
            [time(12, 0)],
265
        ],
266
    )
267
    assert_series_equal(result, expected)
268

269
Product

Resources

Company
