def test_str_split() -> None:
a = pl.Series("a", ["a, b", "a", "ab,c,de"])
for out in [a.str.split(","), pl.select(pl.lit(a).str.split(",")).to_series()]:
assert out[0].to_list() == ["a", " b"]
assert out[1].to_list() == ["a"]
assert out[2].to_list() == ["ab", "c", "de"]
for out in [
a.str.split(",", inclusive=True),
pl.select(pl.lit(a).str.split(",", inclusive=True)).to_series(),
]:
assert out[0].to_list() == ["a,", " b"]
assert out[1].to_list() == ["a"]
assert out[2].to_list() == ["ab,", "c,", "de"]
def test_shuffle() -> None:
a = pl.Series("a", [1, 2, 3])
out = a.shuffle(2)
expected = pl.Series("a", [2, 1, 3])
testing.assert_series_equal(out, expected)
out = pl.select(pl.lit(a).shuffle(2)).to_series()
testing.assert_series_equal(out, expected)
def test_contains() -> None:
a = pl.Series("a", [[1, 2, 3], [2, 5], [6, 7, 8, 9]])
out = a.arr.contains(2)
expected = pl.Series("a", [True, True, False])
testing.assert_series_equal(out, expected)
out = pl.select(pl.lit(a).arr.contains(2)).to_series()
testing.assert_series_equal(out, expected)
def test_list_arr_get() -> None:
a = pl.Series("a", [[1, 2, 3], [4, 5], [6, 7, 8, 9]])
out = a.arr.get(0)
expected = pl.Series("a", [1, 4, 6])
testing.assert_series_equal(out, expected)
out = a.arr.first()
testing.assert_series_equal(out, expected)
out = pl.select(pl.lit(a).arr.first()).to_series()
testing.assert_series_equal(out, expected)
out = a.arr.get(-1)
expected = pl.Series("a", [3, 5, 9])
testing.assert_series_equal(out, expected)
out = a.arr.last()
testing.assert_series_equal(out, expected)
out = pl.select(pl.lit(a).arr.last()).to_series()
testing.assert_series_equal(out, expected)
a = pl.Series("a", [[1, 2, 3], [4, 5], [6, 7, 8, 9]])
out = a.arr.get(-3)
expected = pl.Series("a", [1, None, 7])
testing.assert_series_equal(out, expected)
def test_abs() -> None:
# ints
s = pl.Series([1, -2, 3, -4])
testing.assert_series_equal(s.abs(), pl.Series([1, 2, 3, 4]))
testing.assert_series_equal(np.abs(s), pl.Series([1, 2, 3, 4])) # type: ignore
# floats
s = pl.Series([1.0, -2.0, 3, -4.0])
testing.assert_series_equal(s.abs(), pl.Series([1.0, 2.0, 3.0, 4.0]))
testing.assert_series_equal(
np.abs(s), pl.Series([1.0, 2.0, 3.0, 4.0]) # type: ignore
)
testing.assert_series_equal(
pl.select(pl.lit(s).abs()).to_series(), pl.Series([1.0, 2.0, 3.0, 4.0])
)
def test_list_arr_get() -> None:
a = pl.Series("a", [[1, 2, 3], [4, 5], [6, 7, 8, 9]])
out = a.arr.get(0)
expected = pl.Series("a", [1, 4, 6])
assert_series_equal(out, expected)
out = a.arr.first()
assert_series_equal(out, expected)
out = pl.select(pl.lit(a).arr.first()).to_series()
assert_series_equal(out, expected)
out = a.arr.get(-1)
expected = pl.Series("a", [3, 5, 9])
assert_series_equal(out, expected)
out = a.arr.last()
assert_series_equal(out, expected)
out = pl.select(pl.lit(a).arr.last()).to_series()
assert_series_equal(out, expected)
a = pl.Series("a", [[1, 2, 3], [4, 5], [6, 7, 8, 9]])
out = a.arr.get(-3)
expected = pl.Series("a", [1, None, 7])
assert_series_equal(out, expected)
assert pl.DataFrame(
{"a": [[1], [2], [3], [4, 5, 6], [7, 8, 9], [None, 11]]}
).with_columns(
[pl.col("a").arr.get(i).alias(f"get_{i}") for i in range(4)]
).to_dict(
False
) == {
"a": [[1], [2], [3], [4, 5, 6], [7, 8, 9], [None, 11]],
"get_0": [1, 2, 3, 4, 7, None],
"get_1": [None, None, None, 5, 8, 11],
"get_2": [None, None, None, 6, 9, None],
"get_3": [None, None, None, None, None, None],
}
def test_reshape() -> None:
s = pl.Series("a", [1, 2, 3, 4])
out = s.reshape((-1, 2))
expected = pl.Series("a", [[1, 2], [3, 4]])
assert out.series_equal(expected)
out = s.reshape((2, 2))
assert out.series_equal(expected)
out = s.reshape((2, -1))
assert out.series_equal(expected)
out = s.reshape((-1, 1))
expected = pl.Series("a", [[1], [2], [3], [4]])
assert out.series_equal(expected)
# test lazy_dispatch
out = pl.select(pl.lit(s).reshape((-1, 1))).to_series()
assert out.series_equal(expected)
def test_numpy_to_lit() -> None:
out = pl.select(pl.lit(np.array([1, 2, 3]))).to_series().to_list()
assert out == [1, 2, 3]
out = pl.select(pl.lit(np.float32(0))).to_series().to_list()
assert out == [0.0]
def test_sample() -> None:
a = pl.Series("a", range(0, 20))
out = pl.select(pl.lit(a).sample(0.5, False, seed=1)).to_series()
assert out.shape == (10, )
assert out.to_list() != out.sort().to_list()
assert out.unique().shape == (10, )