From 54e87a50355b251d5bd23b6dd2c81dcb334f5a43 Mon Sep 17 00:00:00 2001 From: barak1412 Date: Fri, 27 Sep 2024 17:49:36 +0300 Subject: [PATCH] added test --- .../tests/unit/functions/test_when_then.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/py-polars/tests/unit/functions/test_when_then.py b/py-polars/tests/unit/functions/test_when_then.py index 9d0b5ecaced0..b84b42dd22f5 100644 --- a/py-polars/tests/unit/functions/test_when_then.py +++ b/py-polars/tests/unit/functions/test_when_then.py @@ -597,6 +597,35 @@ def test_when_then_parametric( assert ref["if_true"].to_list() == ans["if_true"].to_list() +def test_when_then_else_struct_18961() -> None: + v1 = [None, {"foo": 0, "bar": "1"}] + v2 = [{"foo": 0, "bar": "1"}, {"foo": 0, "bar": "1"}] + + df = pl.DataFrame({"left": v1, "right": v2, "mask": [False, True]}) + + expected = [{"foo": 0, "bar": "1"}, {"foo": 0, "bar": "1"}] + ans = ( + df.select( + pl.when(pl.col.mask).then(pl.col.left).otherwise(pl.col.right.first()) + ) + .get_column("left") + .to_list() + ) + assert expected == ans + + df = pl.DataFrame({"left": v2, "right": v1, "mask": [True, False]}) + + expected = [{"foo": 0, "bar": "1"}, {"foo": 0, "bar": "1"}] + ans = ( + df.select( + pl.when(pl.col.mask).then(pl.col.left.first()).otherwise(pl.col.right) + ) + .get_column("left") + .to_list() + ) + assert expected == ans + + def test_when_then_supertype_15975() -> None: df = pl.DataFrame({"a": [1, 2, 3]})