def test_panic_exception() -> None:
with pytest.raises(
pl.PanicException,
match=
r"""this operation is not implemented/valid for this dtype: .*""",
):
pl.struct(pl.Series("a", [1, 2, 3]), eager=True).sort()
def test_eager_struct() -> None:
with pytest.raises(pl.DuplicateError,
match="multiple fields with name '' found"):
s = pl.struct([pl.Series([1, 2, 3]),
pl.Series(["a", "b", "c"])],
eager=True)
s = pl.struct([pl.Series("a", [1, 2, 3]),
pl.Series("b", ["a", "b", "c"])],
eager=True)
assert s.dtype == pl.Struct
def test_struct_function_expansion() -> None:
df = pl.DataFrame({
"a": [1, 2, 3, 4],
"b": ["one", "two", "three", "four"],
"c": [9, 8, 7, 6]
})
assert df.with_column(pl.struct(pl.col(["a",
"b"])))["a"].struct.fields == [
"a",
"b",
]
def test_struct_to_list() -> None:
assert pl.DataFrame({
"int": [1, 2],
"str": ["a", "b"],
"bool": [True, None],
"list": [[1, 2], [3]]
}).select([pl.struct(pl.all()).alias("my_struct")
]).to_series().to_list() == [
(1, "a", True, pl.Series([1, 2])),
(2, "b", None, pl.Series([3])),
]
def test_apply_numpy_int_out() -> None:
df = pl.DataFrame({"col1": [2, 4, 8, 16]})
assert df.with_column(
pl.col("col1").apply(lambda x: np.left_shift(x, 8)).alias("result")
).frame_equal(
pl.DataFrame({"col1": [2, 4, 8, 16], "result": [512, 1024, 2048, 4096]})
)
df = pl.DataFrame({"col1": [2, 4, 8, 16], "shift": [1, 1, 2, 2]})
assert df.select(
pl.struct(["col1", "shift"])
.apply(lambda cols: np.left_shift(cols["col1"], cols["shift"]))
.alias("result")
).frame_equal(pl.DataFrame({"result": [4, 8, 32, 64]}))
def test_apply_struct() -> None:
df = pl.DataFrame({
"A": ["a", "a"],
"B": [2, 3],
"C": [True, False],
"D": [12.0, None]
})
out = df.with_column(pl.struct(df.columns).alias("struct")).select([
pl.col("struct").apply(lambda x: x["A"]).alias("A_field"),
pl.col("struct").apply(lambda x: x["B"]).alias("B_field"),
pl.col("struct").apply(lambda x: x["C"]).alias("C_field"),
pl.col("struct").apply(lambda x: x["D"]).alias("D_field"),
])
expected = pl.DataFrame({
"A_field": ["a", "a"],
"B_field": [2, 3],
"C_field": [True, False],
"D_field": [12.0, None],
})
assert out.frame_equal(expected)
def test_struct_to_list() -> None:
assert pl.DataFrame({
"int": [1, 2],
"str": ["a", "b"],
"bool": [True, None],
"list": [[1, 2], [3]]
}).select([pl.struct(pl.all()).alias("my_struct")
]).to_series().to_list() == [
{
"int": 1,
"str": "a",
"bool": True,
"list": [1, 2]
},
{
"int": 2,
"str": "b",
"bool": None,
"list": [3]
},
]
def map_expr(name: str) -> pl.Expr:
return (pl.when(ignore_nulls or pl.col(name).null_count() == 0).then(
pl.struct([
pl.sum(name).alias("sum"),
(pl.count() - pl.col(name).null_count()).alias("count"),
]), ).otherwise(None)).alias("out")
