CoCalc -- test

GitHub Repository: pola-rs/polars
Path: blob/main/py-polars/tests/unit/operations/namespaces/array/test_array.py
⁶⁹⁴⁰ views
1
from __future__ import annotations
2

3
import datetime
4
from typing import Any
5

6
import pytest
7

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

12

13
def test_arr_min_max() -> None:
14
    s = pl.Series("a", [[1, 2], [4, 3]], dtype=pl.Array(pl.Int64, 2))
15
    assert s.arr.max().to_list() == [2, 4]
16
    assert s.arr.min().to_list() == [1, 3]
17

18
    s_with_null = pl.Series("a", [[None, 2], None, [3, 4]], dtype=pl.Array(pl.Int64, 2))
19
    assert s_with_null.arr.max().to_list() == [2, None, 4]
20
    assert s_with_null.arr.min().to_list() == [2, None, 3]
21

22

23
def test_array_min_max_dtype_12123() -> None:
24
    df = pl.LazyFrame(
25
        [pl.Series("a", [[1.0, 3.0], [2.0, 5.0]]), pl.Series("b", [1.0, 2.0])],
26
        schema_overrides={
27
            "a": pl.Array(pl.Float64, 2),
28
        },
29
    )
30

31
    df = df.with_columns(
32
        max=pl.col("a").arr.max().alias("max"),
33
        min=pl.col("a").arr.min().alias("min"),
34
    )
35

36
    assert df.collect_schema() == {
37
        "a": pl.Array(pl.Float64, 2),
38
        "b": pl.Float64,
39
        "max": pl.Float64,
40
        "min": pl.Float64,
41
    }
42

43
    out = df.select(pl.col("max") * pl.col("b"), pl.col("min") * pl.col("b")).collect()
44

45
    assert_frame_equal(out, pl.DataFrame({"max": [3.0, 10.0], "min": [1.0, 4.0]}))
46

47

48
@pytest.mark.parametrize(
49
    ("data", "expected_sum", "dtype"),
50
    [
51
        ([[1, 2], [4, 3]], [3, 7], pl.Int64),
52
        ([[1, None], [None, 3], [None, None]], [1, 3, 0], pl.Int64),
53
        ([[1.0, 2.0], [4.0, 3.0]], [3.0, 7.0], pl.Float32),
54
        ([[1.0, None], [None, 3.0], [None, None]], [1.0, 3.0, 0], pl.Float32),
55
        ([[True, False], [True, True], [False, False]], [1, 2, 0], pl.Boolean),
56
        ([[True, None], [None, False], [None, None]], [1, 0, 0], pl.Boolean),
57
    ],
58
)
59
def test_arr_sum(
60
    data: list[list[Any]], expected_sum: list[Any], dtype: pl.DataType
61
) -> None:
62
    s = pl.Series("a", data, dtype=pl.Array(dtype, 2))
63
    assert s.arr.sum().to_list() == expected_sum
64

65

66
@pytest.mark.may_fail_cloud
67
def test_array_lengths_zwa() -> None:
68
    assert pl.Series("a", [[], []], pl.Array(pl.Null, 0)).arr.len().to_list() == [0, 0]
69
    assert pl.Series("a", [None, []], pl.Array(pl.Null, 0)).arr.len().to_list() == [
70
        None,
71
        0,
72
    ]
73
    assert pl.Series("a", [None], pl.Array(pl.Null, 0)).arr.len().to_list() == [None]
74

75
    assert pl.Series("a", [], pl.Array(pl.Null, 0)).arr.len().to_list() == []
76

77

78
def test_array_lengths() -> None:
79
    df = pl.DataFrame(
80
        [
81
            pl.Series("a", [[1, 2, 3]], dtype=pl.Array(pl.Int64, 3)),
82
            pl.Series("b", [[4, 5]], dtype=pl.Array(pl.Int64, 2)),
83
        ]
84
    )
85
    out = df.select(pl.col("a").arr.len(), pl.col("b").arr.len())
86
    expected_df = pl.DataFrame(
87
        {"a": [3], "b": [2]}, schema={"a": pl.UInt32, "b": pl.UInt32}
88
    )
89
    assert_frame_equal(out, expected_df)
90

91
    assert pl.Series("a", [], pl.Array(pl.Null, 1)).arr.len().to_list() == []
92
    assert pl.Series(
93
        "a", [[1, 2, 3], None, [7, 8, 9]], pl.Array(pl.Int32, 3)
94
    ).arr.len().to_list() == [3, None, 3]
95

96

97
@pytest.mark.parametrize(
98
    ("as_array"),
99
    [True, False],
100
)
101
def test_arr_slice(as_array: bool) -> None:
102
    df = pl.DataFrame(
103
        {
104
            "arr": [[1, 2, 3], [10, 2, 1]],
105
        },
106
        schema={"arr": pl.Array(pl.Int64, 3)},
107
    )
108

109
    assert df.select([pl.col("arr").arr.slice(0, 1, as_array=as_array)]).to_dict(
110
        as_series=False
111
    ) == {"arr": [[1], [10]]}
112
    assert df.select([pl.col("arr").arr.slice(1, 1, as_array=as_array)]).to_dict(
113
        as_series=False
114
    ) == {"arr": [[2], [2]]}
115
    assert df.select([pl.col("arr").arr.slice(-1, 1, as_array=as_array)]).to_dict(
116
        as_series=False
117
    ) == {"arr": [[3], [1]]}
118
    assert df.select([pl.col("arr").arr.slice(-2, 1, as_array=as_array)]).to_dict(
119
        as_series=False
120
    ) == {"arr": [[2], [2]]}
121
    assert df.select([pl.col("arr").arr.slice(-2, 2, as_array=as_array)]).to_dict(
122
        as_series=False
123
    ) == {"arr": [[2, 3], [2, 1]]}
124
    return
125

126

127
@pytest.mark.parametrize(
128
    ("as_array"),
129
    [True, False],
130
)
131
def test_arr_slice_on_series(as_array: bool) -> None:
132
    vals = [[1, 2, 3, 4], [10, 2, 1, 2]]
133
    s = pl.Series("a", vals, dtype=pl.Array(pl.Int64, 4))
134
    assert s.arr.head(2, as_array=as_array).to_list() == [[1, 2], [10, 2]]
135
    assert s.arr.tail(2, as_array=as_array).to_list() == [[3, 4], [1, 2]]
136
    assert s.arr.tail(10, as_array=as_array).to_list() == vals
137
    assert s.arr.head(10, as_array=as_array).to_list() == vals
138
    assert s.arr.slice(1, 2, as_array=as_array).to_list() == [[2, 3], [2, 1]]
139
    assert s.arr.slice(-5, 2, as_array=as_array).to_list() == [[1], [10]]
140
    # TODO: there is a bug in list.slice that does not allow negative values for head
141
    if as_array:
142
        assert s.arr.tail(-1, as_array=as_array).to_list() == [[2, 3, 4], [2, 1, 2]]
143
        assert s.arr.tail(-2, as_array=as_array).to_list() == [[3, 4], [1, 2]]
144
        assert s.arr.tail(-3, as_array=as_array).to_list() == [[4], [2]]
145
        assert s.arr.head(-1, as_array=as_array).to_list() == [[1, 2, 3], [10, 2, 1]]
146
        assert s.arr.head(-2, as_array=as_array).to_list() == [[1, 2], [10, 2]]
147
        assert s.arr.head(-3, as_array=as_array).to_list() == [[1], [10]]
148

149

150
def test_arr_unique() -> None:
151
    df = pl.DataFrame(
152
        {"a": pl.Series("a", [[1, 1], [4, 3]], dtype=pl.Array(pl.Int64, 2))}
153
    )
154

155
    out = df.select(pl.col("a").arr.unique(maintain_order=True))
156
    expected = pl.DataFrame({"a": [[1], [4, 3]]})
157
    assert_frame_equal(out, expected)
158

159

160
def test_array_any_all() -> None:
161
    s = pl.Series(
162
        [[True, True], [False, True], [False, False], [None, None], None],
163
        dtype=pl.Array(pl.Boolean, 2),
164
    )
165

166
    expected_any = pl.Series([True, True, False, False, None])
167
    assert_series_equal(s.arr.any(), expected_any)
168

169
    expected_all = pl.Series([True, False, False, True, None])
170
    assert_series_equal(s.arr.all(), expected_all)
171

172
    s = pl.Series([[1, 2], [3, 4], [5, 6]], dtype=pl.Array(pl.Int64, 2))
173
    with pytest.raises(ComputeError, match="expected boolean elements in array"):
174
        s.arr.any()
175
    with pytest.raises(ComputeError, match="expected boolean elements in array"):
176
        s.arr.all()
177

178

179
def test_array_sort() -> None:
180
    s = pl.Series([[2, None, 1], [1, 3, 2]], dtype=pl.Array(pl.UInt32, 3))
181

182
    desc = s.arr.sort(descending=True)
183
    expected = pl.Series([[None, 2, 1], [3, 2, 1]], dtype=pl.Array(pl.UInt32, 3))
184
    assert_series_equal(desc, expected)
185

186
    asc = s.arr.sort(descending=False)
187
    expected = pl.Series([[None, 1, 2], [1, 2, 3]], dtype=pl.Array(pl.UInt32, 3))
188
    assert_series_equal(asc, expected)
189

190
    # test nulls_last
191
    s = pl.Series([[None, 1, 2], [-1, None, 9]], dtype=pl.Array(pl.Int8, 3))
192
    assert_series_equal(
193
        s.arr.sort(nulls_last=True),
194
        pl.Series([[1, 2, None], [-1, 9, None]], dtype=pl.Array(pl.Int8, 3)),
195
    )
196
    assert_series_equal(
197
        s.arr.sort(nulls_last=False),
198
        pl.Series([[None, 1, 2], [None, -1, 9]], dtype=pl.Array(pl.Int8, 3)),
199
    )
200

201

202
def test_array_reverse() -> None:
203
    s = pl.Series([[2, None, 1], [1, None, 2]], dtype=pl.Array(pl.UInt32, 3))
204

205
    s = s.arr.reverse()
206
    expected = pl.Series([[1, None, 2], [2, None, 1]], dtype=pl.Array(pl.UInt32, 3))
207
    assert_series_equal(s, expected)
208

209

210
def test_array_arg_min_max() -> None:
211
    s = pl.Series("a", [[1, 2, 4], [3, 2, 1]], dtype=pl.Array(pl.UInt32, 3))
212
    expected = pl.Series("a", [0, 2], dtype=pl.UInt32)
213
    assert_series_equal(s.arr.arg_min(), expected)
214
    expected = pl.Series("a", [2, 0], dtype=pl.UInt32)
215
    assert_series_equal(s.arr.arg_max(), expected)
216

217

218
def test_array_get() -> None:
219
    s = pl.Series(
220
        "a",
221
        [[1, 2, 3, 4], [5, 6, None, None], [7, 8, 9, 10]],
222
        dtype=pl.Array(pl.Int64, 4),
223
    )
224

225
    # Test index literal.
226
    out = s.arr.get(1, null_on_oob=False)
227
    expected = pl.Series("a", [2, 6, 8], dtype=pl.Int64)
228
    assert_series_equal(out, expected)
229

230
    # Null index literal.
231
    out_df = s.to_frame().select(pl.col.a.arr.get(pl.lit(None), null_on_oob=False))
232
    expected_df = pl.Series("a", [None, None, None], dtype=pl.Int64).to_frame()
233
    assert_frame_equal(out_df, expected_df)
234

235
    # Out-of-bounds index literal.
236
    with pytest.raises(ComputeError, match="get index is out of bounds"):
237
        out = s.arr.get(100, null_on_oob=False)
238

239
    # Negative index literal.
240
    out = s.arr.get(-2, null_on_oob=False)
241
    expected = pl.Series("a", [3, None, 9], dtype=pl.Int64)
242
    assert_series_equal(out, expected)
243

244
    # Test index expr.
245
    with pytest.raises(ComputeError, match="get index is out of bounds"):
246
        out = s.arr.get(pl.Series([1, -2, 100]), null_on_oob=False)
247

248
    out = s.arr.get(pl.Series([1, -2, 0]), null_on_oob=False)
249
    expected = pl.Series("a", [2, None, 7], dtype=pl.Int64)
250
    assert_series_equal(out, expected)
251

252
    # Test logical type.
253
    s = pl.Series(
254
        "a",
255
        [
256
            [datetime.date(1999, 1, 1), datetime.date(2000, 1, 1)],
257
            [datetime.date(2001, 10, 1), None],
258
            [None, None],
259
        ],
260
        dtype=pl.Array(pl.Date, 2),
261
    )
262
    with pytest.raises(ComputeError, match="get index is out of bounds"):
263
        out = s.arr.get(pl.Series([1, -2, 4]), null_on_oob=False)
264

265

266
def test_array_get_null_on_oob() -> None:
267
    s = pl.Series(
268
        "a",
269
        [[1, 2, 3, 4], [5, 6, None, None], [7, 8, 9, 10]],
270
        dtype=pl.Array(pl.Int64, 4),
271
    )
272

273
    # Test index literal.
274
    out = s.arr.get(1, null_on_oob=True)
275
    expected = pl.Series("a", [2, 6, 8], dtype=pl.Int64)
276
    assert_series_equal(out, expected)
277

278
    # Null index literal.
279
    out_df = s.to_frame().select(pl.col.a.arr.get(pl.lit(None), null_on_oob=True))
280
    expected_df = pl.Series("a", [None, None, None], dtype=pl.Int64).to_frame()
281
    assert_frame_equal(out_df, expected_df)
282

283
    # Out-of-bounds index literal.
284
    out = s.arr.get(100, null_on_oob=True)
285
    expected = pl.Series("a", [None, None, None], dtype=pl.Int64)
286
    assert_series_equal(out, expected)
287

288
    # Negative index literal.
289
    out = s.arr.get(-2, null_on_oob=True)
290
    expected = pl.Series("a", [3, None, 9], dtype=pl.Int64)
291
    assert_series_equal(out, expected)
292

293
    # Test index expr.
294
    out = s.arr.get(pl.Series([1, -2, 100]), null_on_oob=True)
295
    expected = pl.Series("a", [2, None, None], dtype=pl.Int64)
296
    assert_series_equal(out, expected)
297

298
    # Test logical type.
299
    s = pl.Series(
300
        "a",
301
        [
302
            [datetime.date(1999, 1, 1), datetime.date(2000, 1, 1)],
303
            [datetime.date(2001, 10, 1), None],
304
            [None, None],
305
        ],
306
        dtype=pl.Array(pl.Date, 2),
307
    )
308
    out = s.arr.get(pl.Series([1, -2, 4]), null_on_oob=True)
309
    expected = pl.Series(
310
        "a",
311
        [datetime.date(2000, 1, 1), datetime.date(2001, 10, 1), None],
312
        dtype=pl.Date,
313
    )
314
    assert_series_equal(out, expected)
315

316

317
def test_arr_first_last() -> None:
318
    s = pl.Series(
319
        "a",
320
        [[1, 2, 3], [None, 5, 6], [None, None, None]],
321
        dtype=pl.Array(pl.Int64, 3),
322
    )
323

324
    first = s.arr.first()
325
    expected_first = pl.Series(
326
        "a",
327
        [1, None, None],
328
        dtype=pl.Int64,
329
    )
330
    assert_series_equal(first, expected_first)
331

332
    last = s.arr.last()
333
    expected_last = pl.Series(
334
        "a",
335
        [3, 6, None],
336
        dtype=pl.Int64,
337
    )
338
    assert_series_equal(last, expected_last)
339

340

341
@pytest.mark.parametrize(
342
    ("data", "set", "dtype"),
343
    [
344
        ([1, 2], [[1, 2], [3, 4]], pl.Int64),
345
        ([True, False], [[True, False], [True, True]], pl.Boolean),
346
        (["a", "b"], [["a", "b"], ["c", "d"]], pl.String),
347
        ([b"a", b"b"], [[b"a", b"b"], [b"c", b"d"]], pl.Binary),
348
        (
349
            [{"a": 1}, {"a": 2}],
350
            [[{"a": 1}, {"a": 2}], [{"b": 1}, {"a": 3}]],
351
            pl.Struct([pl.Field("a", pl.Int64)]),
352
        ),
353
    ],
354
)
355
def test_is_in_array(data: list[Any], set: list[list[Any]], dtype: pl.DataType) -> None:
356
    df = pl.DataFrame(
357
        {"a": data, "arr": set},
358
        schema={"a": dtype, "arr": pl.Array(dtype, 2)},
359
    )
360
    out = df.select(is_in=pl.col("a").is_in(pl.col("arr"))).to_series()
361
    expected = pl.Series("is_in", [True, False])
362
    assert_series_equal(out, expected)
363

364

365
def test_array_join() -> None:
366
    df = pl.DataFrame(
367
        {
368
            "a": [["ab", "c", "d"], ["e", "f", "g"], [None, None, None], None],
369
            "separator": ["&", None, "*", "_"],
370
        },
371
        schema={
372
            "a": pl.Array(pl.String, 3),
373
            "separator": pl.String,
374
        },
375
    )
376
    out = df.select(pl.col("a").arr.join("-"))
377
    assert out.to_dict(as_series=False) == {"a": ["ab-c-d", "e-f-g", "", None]}
378
    out = df.select(pl.col("a").arr.join(pl.col("separator")))
379
    assert out.to_dict(as_series=False) == {"a": ["ab&c&d", None, "", None]}
380

381
    # test ignore_nulls argument
382
    df = pl.DataFrame(
383
        {
384
            "a": [
385
                ["a", None, "b", None],
386
                None,
387
                [None, None, None, None],
388
                ["c", "d", "e", "f"],
389
            ],
390
            "separator": ["-", "&", " ", "@"],
391
        },
392
        schema={
393
            "a": pl.Array(pl.String, 4),
394
            "separator": pl.String,
395
        },
396
    )
397
    # ignore nulls
398
    out = df.select(pl.col("a").arr.join("-", ignore_nulls=True))
399
    assert out.to_dict(as_series=False) == {"a": ["a-b", None, "", "c-d-e-f"]}
400
    out = df.select(pl.col("a").arr.join(pl.col("separator"), ignore_nulls=True))
401
    assert out.to_dict(as_series=False) == {"a": ["a-b", None, "", "c@d@e@f"]}
402
    # propagate nulls
403
    out = df.select(pl.col("a").arr.join("-", ignore_nulls=False))
404
    assert out.to_dict(as_series=False) == {"a": [None, None, None, "c-d-e-f"]}
405
    out = df.select(pl.col("a").arr.join(pl.col("separator"), ignore_nulls=False))
406
    assert out.to_dict(as_series=False) == {"a": [None, None, None, "c@d@e@f"]}
407

408

409
def test_array_explode() -> None:
410
    df = pl.DataFrame(
411
        {
412
            "str": [["a", "b"], ["c", None], None],
413
            "nested": [[[1, 2], [3]], [[], [4, None]], None],
414
            "logical": [
415
                [datetime.date(1998, 1, 1), datetime.date(2000, 10, 1)],
416
                [datetime.date(2024, 1, 1), None],
417
                None,
418
            ],
419
        },
420
        schema={
421
            "str": pl.Array(pl.String, 2),
422
            "nested": pl.Array(pl.List(pl.Int64), 2),
423
            "logical": pl.Array(pl.Date, 2),
424
        },
425
    )
426
    out = df.select(pl.all().arr.explode())
427
    expected = pl.DataFrame(
428
        {
429
            "str": ["a", "b", "c", None, None],
430
            "nested": [[1, 2], [3], [], [4, None], None],
431
            "logical": [
432
                datetime.date(1998, 1, 1),
433
                datetime.date(2000, 10, 1),
434
                datetime.date(2024, 1, 1),
435
                None,
436
                None,
437
            ],
438
        }
439
    )
440
    assert_frame_equal(out, expected)
441

442
    # test no-null fast path
443
    s = pl.Series(
444
        [
445
            [datetime.date(1998, 1, 1), datetime.date(1999, 1, 3)],
446
            [datetime.date(2000, 1, 1), datetime.date(2023, 10, 1)],
447
        ],
448
        dtype=pl.Array(pl.Date, 2),
449
    )
450
    out_s = s.arr.explode()
451
    expected_s = pl.Series(
452
        [
453
            datetime.date(1998, 1, 1),
454
            datetime.date(1999, 1, 3),
455
            datetime.date(2000, 1, 1),
456
            datetime.date(2023, 10, 1),
457
        ],
458
        dtype=pl.Date,
459
    )
460
    assert_series_equal(out_s, expected_s)
461

462

463
@pytest.mark.parametrize(
464
    ("arr", "data", "expected", "dtype"),
465
    [
466
        ([[1, 2], [3, None], None], 1, [1, 0, None], pl.Int64),
467
        ([[True, False], [True, None], None], True, [1, 1, None], pl.Boolean),
468
        ([["a", "b"], ["c", None], None], "a", [1, 0, None], pl.String),
469
        ([[b"a", b"b"], [b"c", None], None], b"a", [1, 0, None], pl.Binary),
470
    ],
471
)
472
def test_array_count_matches(
473
    arr: list[list[Any] | None], data: Any, expected: list[Any], dtype: pl.DataType
474
) -> None:
475
    df = pl.DataFrame({"arr": arr}, schema={"arr": pl.Array(dtype, 2)})
476
    out = df.select(count_matches=pl.col("arr").arr.count_matches(data))
477
    assert out.to_dict(as_series=False) == {"count_matches": expected}
478

479

480
def test_array_count_matches_wildcard_expansion() -> None:
481
    df = pl.DataFrame(
482
        {"a": [[1, 2]], "b": [[3, 4]]},
483
        schema={"a": pl.Array(pl.Int64, 2), "b": pl.Array(pl.Int64, 2)},
484
    )
485
    assert df.select(pl.all().arr.count_matches(3)).to_dict(as_series=False) == {
486
        "a": [0],
487
        "b": [1],
488
    }
489

490

491
def test_array_to_struct() -> None:
492
    df = pl.DataFrame(
493
        {"a": [[1, 2, 3], [4, 5, None]]}, schema={"a": pl.Array(pl.Int8, 3)}
494
    )
495
    assert df.select([pl.col("a").arr.to_struct()]).to_series().to_list() == [
496
        {"field_0": 1, "field_1": 2, "field_2": 3},
497
        {"field_0": 4, "field_1": 5, "field_2": None},
498
    ]
499

500
    df = pl.DataFrame(
501
        {"a": [[1, 2, None], [1, 2, 3]]}, schema={"a": pl.Array(pl.Int8, 3)}
502
    )
503
    assert df.select(
504
        pl.col("a").arr.to_struct(fields=lambda idx: f"col_name_{idx}")
505
    ).to_series().to_list() == [
506
        {"col_name_0": 1, "col_name_1": 2, "col_name_2": None},
507
        {"col_name_0": 1, "col_name_1": 2, "col_name_2": 3},
508
    ]
509

510
    assert df.lazy().select(pl.col("a").arr.to_struct()).unnest(
511
        "a"
512
    ).sum().collect().columns == ["field_0", "field_1", "field_2"]
513

514

515
def test_array_shift() -> None:
516
    df = pl.DataFrame(
517
        {"a": [[1, 2, 3], None, [4, 5, 6], [7, 8, 9]], "n": [None, 1, 1, -2]},
518
        schema={"a": pl.Array(pl.Int64, 3), "n": pl.Int64},
519
    )
520

521
    out = df.select(
522
        lit=pl.col("a").arr.shift(1), expr=pl.col("a").arr.shift(pl.col("n"))
523
    )
524
    expected = pl.DataFrame(
525
        {
526
            "lit": [[None, 1, 2], None, [None, 4, 5], [None, 7, 8]],
527
            "expr": [None, None, [None, 4, 5], [9, None, None]],
528
        },
529
        schema={"lit": pl.Array(pl.Int64, 3), "expr": pl.Array(pl.Int64, 3)},
530
    )
531
    assert_frame_equal(out, expected)
532

533

534
def test_array_n_unique() -> None:
535
    df = pl.DataFrame(
536
        {
537
            "a": [[1, 1, 2], [3, 3, 3], [None, None, None], None],
538
        },
539
        schema={"a": pl.Array(pl.Int64, 3)},
540
    )
541

542
    out = df.select(n_unique=pl.col("a").arr.n_unique())
543
    expected = pl.DataFrame(
544
        {"n_unique": [2, 1, 1, None]}, schema={"n_unique": pl.UInt32}
545
    )
546
    assert_frame_equal(out, expected)
547

548

549
def test_explode_19049() -> None:
550
    df = pl.DataFrame({"a": [[1, 2, 3]]}, schema={"a": pl.Array(pl.Int64, 3)})
551
    result_df = df.select(pl.col.a.arr.explode())
552
    expected_df = pl.DataFrame({"a": [1, 2, 3]}, schema={"a": pl.Int64})
553
    assert_frame_equal(result_df, expected_df)
554

555
    df = pl.DataFrame({"a": [1, 2, 3]}, schema={"a": pl.Int64})
556
    with pytest.raises(InvalidOperationError, match="expected Array type, got: i64"):
557
        df.select(pl.col.a.arr.explode())
558

559

560
def test_array_join_unequal_lengths_22018() -> None:
561
    df = pl.DataFrame(
562
        [
563
            pl.Series(
564
                "a",
565
                [
566
                    ["a", "b", "d"],
567
                    ["ya", "x", "y"],
568
                    ["ya", "x", "y"],
569
                ],
570
                pl.Array(pl.String, 3),
571
            ),
572
        ]
573
    )
574
    with pytest.raises(pl.exceptions.ShapeError):
575
        df.select(pl.col.a.arr.join(pl.Series([",", "-"])))
576

577

578
def test_array_shift_unequal_lengths_22018() -> None:
579
    with pytest.raises(pl.exceptions.ShapeError):
580
        pl.Series(
581
            "a",
582
            [
583
                ["a", "b", "d"],
584
                ["a", "b", "d"],
585
                ["a", "b", "d"],
586
            ],
587
            pl.Array(pl.String, 3),
588
        ).arr.shift(pl.Series([1, 2]))
589

590

591
def test_array_shift_self_broadcast_22124() -> None:
592
    assert_series_equal(
593
        pl.Series(
594
            "a",
595
            [
596
                ["a", "b", "d"],
597
            ],
598
            pl.Array(pl.String, 3),
599
        ).arr.shift(pl.Series([1, 2])),
600
        pl.Series(
601
            "a",
602
            [
603
                [None, "a", "b"],
604
                [None, None, "a"],
605
            ],
606
            pl.Array(pl.String, 3),
607
        ),
608
    )
609

610

611
def test_arr_contains() -> None:
612
    s = pl.Series([[1, 2, None], [None, None, None], None], dtype=pl.Array(pl.Int64, 3))
613

614
    assert_series_equal(
615
        s.arr.contains(None, nulls_equal=False),
616
        pl.Series([None, None, None], dtype=pl.Boolean),
617
    )
618
    assert_series_equal(
619
        s.arr.contains(None, nulls_equal=True),
620
        pl.Series([True, True, None], dtype=pl.Boolean),
621
    )
622
    assert_series_equal(
623
        s.arr.contains(1, nulls_equal=False),
624
        pl.Series([True, False, None], dtype=pl.Boolean),
625
    )
626
    assert_series_equal(
627
        s.arr.contains(1, nulls_equal=True),
628
        pl.Series([True, False, None], dtype=pl.Boolean),
629
    )
630

631
Product

Resources

Company
