def test_repeat(): s = pl.repeat(1, 10) assert s.dtype == pl.Int64 assert s.len() == 10 s = pl.repeat("foo", 10) assert s.dtype == pl.Utf8 assert s.len() == 10
def test_expression_appends() -> None: df = pl.DataFrame({"a": [1, 1, 2]}) assert df.select(pl.repeat(None, 3).append(pl.col("a"))).n_chunks() == 2 assert df.select(pl.repeat(None, 3).append(pl.col("a")).rechunk()).n_chunks() == 1 out = df.select(pl.concat([pl.repeat(None, 3), pl.col("a")])) assert out.n_chunks() == 1 assert out.to_series().to_list() == [None, None, None, 1, 1, 2]
def test_repeat() -> None: s = pl.repeat(1, 10, eager=True) assert s.dtype == pl.Int64 assert s.len() == 10 s = pl.repeat("foo", 10, eager=True) assert s.dtype == pl.Utf8 assert s.len() == 10 s = pl.repeat(1.0, 5, eager=True) assert s.dtype == pl.Float64 assert s.len() == 5 assert s == [1.0, 1.0, 1.0, 1.0, 1.0] s = pl.repeat(True, 5, eager=True) assert s.dtype == pl.Boolean assert s.len() == 5
def test_repeat_expansion_in_groupby() -> None: out = ( pl.DataFrame({"g": [1, 2, 2, 3, 3, 3]}) .groupby("g", maintain_order=True) .agg(pl.repeat(1, pl.count()).cumsum()) .to_dict() ) assert out == {"g": [1, 2, 3], "literal": [[1], [1, 2], [1, 2, 3]]}
# may contain many things that seemed to go wrong at scale import time import numpy as np import polars as pl # https://github.com/pola-rs/polars/issues/1942 t0 = time.time() pl.repeat(float("nan"), 2 << 12).sort() assert (time.time() - t0) < 1 # test mean overflow issues np.random.seed(1) mean = 769.5607652 df = pl.DataFrame(np.random.randint(500, 1040, 5000000), columns=["value"]) assert np.isclose(df.with_column(pl.mean("value"))[0, 0], mean) assert np.isclose( df.with_column(pl.col("value").cast(pl.Int32)).with_column(pl.mean("value"))[0, 0], mean, ) assert np.isclose( df.with_column(pl.col("value").cast(pl.Int32)).get_column("value").mean(), mean ) # https://github.com/pola-rs/polars/issues/2850 df = pl.DataFrame( { "id": [ 130352432,
def test_apply_list_out() -> None: s = pl.Series("count", [3, 2, 2]) out = s.apply(lambda val: pl.repeat(val, val, eager=True)) assert out[0] == [3, 3, 3] assert out[1] == [2, 2] assert out[2] == [2, 2]