diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs index bcf04712a889..b09b6ab4b090 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs @@ -141,13 +141,22 @@ pub(super) fn concat(s: &mut [Series]) -> PolarsResult> { let mut first = std::mem::take(&mut s[0]); let other = &s[1..]; - let first_ca = match first.list().ok() { + let mut first_ca = match first.list().ok() { Some(ca) => ca, None => { first = first.reshape(&[-1, 1]).unwrap(); first.list().unwrap() } - }; + } + .clone(); + + if first_ca.len() == 1 && !other.is_empty() { + let max_len = other.iter().map(|s| s.len()).max().unwrap(); + if max_len > 1 { + first_ca = first_ca.new_from_index(0, max_len) + } + } + first_ca.lst_concat(other).map(|ca| Some(ca.into_series())) } diff --git a/py-polars/tests/unit/datatypes/test_list.py b/py-polars/tests/unit/datatypes/test_list.py index 999747ee01b3..d19825124438 100644 --- a/py-polars/tests/unit/datatypes/test_list.py +++ b/py-polars/tests/unit/datatypes/test_list.py @@ -368,6 +368,18 @@ def test_flat_aggregation_to_list_conversion_6918() -> None: ).to_dict(False) == {"a": [1, 2], "b": [[[0.0, 1.0]], [[3.0, 4.0]]]} +def test_concat_list_with_lit() -> None: + df = pl.DataFrame({"a": [1, 2, 3]}) + + assert df.select(pl.concat_list([pl.col("a"), pl.lit(1)]).alias("a")).to_dict( + False + ) == {"a": [[1, 1], [2, 1], [3, 1]]} + + assert df.select(pl.concat_list([pl.lit(1), pl.col("a")]).alias("a")).to_dict( + False + ) == {"a": [[1, 1], [1, 2], [1, 3]]} + + def test_list_count_match() -> None: assert pl.DataFrame({"listcol": [[], [1], [1, 2, 3, 2], [1, 2, 1], [4, 4]]}).select( pl.col("listcol").arr.count_match(2).alias("number_of_twos")