diff --git a/crates/polars-core/src/frame/explode.rs b/crates/polars-core/src/frame/explode.rs index 77edc1c94f89..774b30f17390 100644 --- a/crates/polars-core/src/frame/explode.rs +++ b/crates/polars-core/src/frame/explode.rs @@ -83,7 +83,18 @@ impl DataFrame { let check_offsets = || { let first_offsets = exploded_columns[0].1.as_slice(); for (_, offsets) in &exploded_columns[1..] { - polars_ensure!(first_offsets == offsets.as_slice(), + let offsets = offsets.as_slice(); + + let offset_l = first_offsets[0]; + let offset_r = offsets[0]; + let all_equal_len = first_offsets.len() != offsets.len() || { + first_offsets + .iter() + .zip(offsets.iter()) + .all(|(l, r)| (*l - offset_l) == (*r - offset_r)) + }; + + polars_ensure!(all_equal_len, ShapeMismatch: "exploded columns must have matching element counts" ) } diff --git a/py-polars/tests/unit/operations/test_explode.py b/py-polars/tests/unit/operations/test_explode.py index fa3e51efabce..c6a5049319a3 100644 --- a/py-polars/tests/unit/operations/test_explode.py +++ b/py-polars/tests/unit/operations/test_explode.py @@ -438,3 +438,12 @@ def test_undefined_col_15852() -> None: with pytest.raises(pl.exceptions.ColumnNotFoundError): lf.explode("bar").join(lf, on="foo").collect() + + +def test_explode_17648() -> None: + df = pl.DataFrame({"a": [[1, 3], [2, 6, 7], [3, 9, 2], [4], [5, 1, 2, 3, 4]]}) + assert ( + df.slice(1, 2) + .with_columns(pl.int_ranges(pl.col("a").list.len()).alias("count")) + .explode("a", "count") + ).to_dict(as_series=False) == {"a": [2, 6, 7, 3, 9, 2], "count": [0, 1, 2, 0, 1, 2]}