Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add for Boolean
Browse files Browse the repository at this point in the history
etiennebacher committed Jan 15, 2025
1 parent 4c97f57 commit 2d38618
Showing 2 changed files with 40 additions and 5 deletions.
15 changes: 15 additions & 0 deletions crates/polars-ops/src/chunked_array/list/namespace.rs
Original file line number Diff line number Diff line change
@@ -590,6 +590,21 @@ pub trait ListNameSpaceImpl: AsList {
Some(filled_values.into())
})
},
DataType::Boolean => {
let fill_value = fill_value.bool()?;
ca.zip_and_apply_amortized(fill_value, |s, fill_value| {
let binding = s.unwrap();
let s: &Series = binding.as_ref();
let ca = s.bool().unwrap();
let mut filled_values = BooleanChunked::new(
PlSmallStr::EMPTY,
vec![fill_value.unwrap(); max_len - ca.len()],
);

let _ = filled_values.append(ca);
Some(filled_values.into())
})
},
dt => {
polars_bail!(InvalidOperation: "list.pad_start() doesn't work on data type {}", dt)
},
30 changes: 25 additions & 5 deletions py-polars/tests/unit/operations/namespaces/list/test_pad.py
Original file line number Diff line number Diff line change
@@ -32,10 +32,31 @@ def test_list_pad_start_with_expr() -> None:


@pytest.mark.may_fail_auto_streaming
def test_list_pad_start_with_lit() -> None:
df = pl.DataFrame({"a": [[1], [], [1, 2, 3]]})
result = df.select(pl.col("a").list.pad_start(fill_value=0))
expected = pl.DataFrame({"a": [[0, 0, 1], [0, 0, 0], [1, 2, 3]]})
@pytest.mark.parametrize(
("data", "fill_value", "expect"),
[
([[1], [], [1, 2, 3]], 0, [[0, 0, 1], [0, 0, 0], [1, 2, 3]]),
(
[[1.0], [], [1.0, 2.0, 3.0]],
0.0,
[[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 3.0]],
),
(
[["a"], [], ["a", "b", "b"]],
"foo",
[["foo", "foo", "a"], ["foo", "foo", "foo"], ["a", "b", "b"]],
),
(
[[True], [], [False, False, True]],
True,
[[True, True, True], [True, True, True], [False, False, True]],
),
],
)
def test_list_pad_start_with_lit(data: Any, fill_value: Any, expect: Any) -> None:
df = pl.DataFrame({"a": data})
result = df.select(pl.col("a").list.pad_start(fill_value))
expected = pl.DataFrame({"a": expect})
assert_frame_equal(result, expected)


@@ -71,7 +92,6 @@ def test_list_pad_start_errors() -> None:
@pytest.mark.parametrize(
("fill_value", "type"),
[
(True, pl.Boolean),
(timedelta(days=1), pl.Duration),
(date(2022, 1, 1), pl.Date),
(datetime(2022, 1, 1, 23), pl.Datetime),

0 comments on commit 2d38618

Please sign in to comment.