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, )