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]
