CoCalc -- test_get

GitHub Repository: pola-rs/polars
Path: blob/main/py-polars/tests/unit/series/buffers/test_get_buffers.py
⁶⁹⁴⁰ views
1
from datetime import date
2
from typing import cast
3

4
import pytest
5

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

9

10
def test_get_buffers_only_values() -> None:
11
    s = pl.Series([1, 2, 3])
12

13
    result = s._get_buffers()
14

15
    assert_series_equal(result["values"], s)
16
    assert result["validity"] is None
17
    assert result["offsets"] is None
18

19

20
def test_get_buffers_with_validity() -> None:
21
    s = pl.Series([1.5, None, 3.5])
22

23
    result = s._get_buffers()
24

25
    expected_values = pl.Series([1.5, 0.0, 3.5])
26
    assert_series_equal(result["values"], expected_values)
27

28
    validity = cast(pl.Series, result["validity"])
29
    expected_validity = pl.Series([True, False, True])
30
    assert_series_equal(validity, expected_validity)
31

32
    assert result["offsets"] is None
33

34

35
def test_get_buffers_string_type() -> None:
36
    s = pl.Series(["a", "bc", None, "éâç", ""])
37

38
    result = s._get_buffers()
39

40
    expected_values = pl.Series(
41
        [97, 98, 99, 195, 169, 195, 162, 195, 167], dtype=pl.UInt8
42
    )
43
    assert_series_equal(result["values"], expected_values)
44

45
    validity = cast(pl.Series, result["validity"])
46
    expected_validity = pl.Series([True, True, False, True, True])
47
    assert_series_equal(validity, expected_validity)
48

49
    offsets = cast(pl.Series, result["offsets"])
50
    expected_offsets = pl.Series([0, 1, 3, 3, 9, 9], dtype=pl.Int64)
51
    assert_series_equal(offsets, expected_offsets)
52

53

54
def test_get_buffers_logical_sliced() -> None:
55
    s = pl.Series([date(1970, 1, 1), None, date(1970, 1, 3)])[1:]
56

57
    result = s._get_buffers()
58

59
    expected_values = pl.Series([0, 2], dtype=pl.Int32)
60
    assert_series_equal(result["values"], expected_values)
61

62
    validity = cast(pl.Series, result["validity"])
63
    expected_validity = pl.Series([False, True])
64
    assert_series_equal(validity, expected_validity)
65

66
    assert result["offsets"] is None
67

68

69
def test_get_buffers_chunked() -> None:
70
    s = pl.Series([1, 2, None, 4], dtype=pl.UInt8)
71
    s_chunked = pl.concat([s[:2], s[2:]], rechunk=False)
72

73
    result = s_chunked._get_buffers()
74

75
    expected_values = pl.Series([1, 2, 0, 4], dtype=pl.UInt8)
76
    assert_series_equal(result["values"], expected_values)
77
    assert result["values"].n_chunks() == 2
78

79
    validity = cast(pl.Series, result["validity"])
80
    expected_validity = pl.Series([True, True, False, True])
81
    assert_series_equal(validity, expected_validity)
82
    assert validity.n_chunks() == 2
83

84

85
def test_get_buffers_chunked_string_type() -> None:
86
    s = pl.Series(["a", "bc", None, "éâç", ""])
87
    s_chunked = pl.concat([s[:2], s[2:]], rechunk=False)
88

89
    result = s_chunked._get_buffers()
90

91
    expected_values = pl.Series(
92
        [97, 98, 99, 195, 169, 195, 162, 195, 167], dtype=pl.UInt8
93
    )
94
    assert_series_equal(result["values"], expected_values)
95
    assert result["values"].n_chunks() == 1
96

97
    validity = cast(pl.Series, result["validity"])
98
    expected_validity = pl.Series([True, True, False, True, True])
99
    assert_series_equal(validity, expected_validity)
100
    assert validity.n_chunks() == 1
101

102
    offsets = cast(pl.Series, result["offsets"])
103
    expected_offsets = pl.Series([0, 1, 3, 3, 9, 9], dtype=pl.Int64)
104
    assert_series_equal(offsets, expected_offsets)
105
    assert offsets.n_chunks() == 1
106

107

108
def test_get_buffers_unsupported_data_type() -> None:
109
    s = pl.Series([[1, 2], [3]])
110

111
    msg = "`_get_buffers` not implemented for `dtype` list\\[i64\\]"
112
    with pytest.raises(TypeError, match=msg):
113
        s._get_buffers()
114

115
Product

Resources

Company
