Calling explode() on multiple columns of a dataframe slice throws an error.

[jackruder]

Checks

• I have checked that this issue has not already been reported.
• I have confirmed this bug exists on the latest version of Polars.

Reproducible example

import polars as pl
df = pl.DataFrame({'a': [[1,3],[2,6,7],[3,9,2],[4],[5,1,2,3,4]]})
df = (
    df
    .slice(1,2)
    .with_columns(
        pl.int_ranges(pl.col('a').list.len()).alias('count')
    )
    .explode('a', 'count')
)
print(df)

Log output

Traceback (most recent call last):
  File "./plreport.py", line 9, in <module>
    .explode('a', 'count')
     ^^^^^^^^^^^^^^^^^^^^^
  File ".virtualenvs/bubbles/lib/python3.12/site-packages/polars/dataframe/frame.py", line 7763, in explode
    return self.lazy().explode(columns, *more_columns).collect(_eager=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".virtualenvs/bubbles/lib/python3.12/site-packages/polars/lazyframe/frame.py", line 1942, in collect
    return wrap_df(ldf.collect(callback))
                   ^^^^^^^^^^^^^^^^^^^^^
polars.exceptions.ShapeError: exploded columns must have matching element counts

Issue description

This issue occurs only when using a slice of a dataframe to perform the computation. For example, the following example runs correctly.

import polars as pl
df = pl.DataFrame({'a': [[1,3],[2,6,7],[3,9,2],[4],[5,1,2,3,4]]})
df = (
    df
    .with_columns(
        pl.int_ranges(pl.col('a').list.len()).alias('count')
    )
    .explode('a', 'count')
)
print(df)

The 'counts' of each list are correct:

print(
    df
    .slice(0,4)
    .with_columns(
         pl.int_ranges(pl.col('a').list.len()).alias('count')
    )
    .with_columns(
        [
            pl.col('a').list.len().alias('a_len'),
            pl.col('count').list.len().alias('count_len')
        ]
    ).select(["a_len", "count_len"])
)

Output (sorry for the rough formatting):
shape: (4, 2)

│ a_len ┆ count_len │
│ u32 ┆ u32 │
╞═══════╪
│ 2 ┆ 2 │
│ 3 ┆ 3 │
│ 3 ┆ 3 │
│ 1 ┆ 1 │
└───────┴

Something like this should also work, but does not, throwing the exact same error. Note that the collect() executes with no issues.

import polars as pl
df = pl.LazyFrame({'a': [[1,3],[2,6,7],[3,9,2],[4],[5,1,2,3,4]]})
df = (
    df
    .slice(1,2)
    .with_columns(
        pl.int_ranges(pl.col('a').list.len()).alias('count')
    )
)
print(df.collect().clone().explode('a', 'count'))

Expected behavior

The example should run, and produce
┌──────┐
│ a ┆ count │
│ i64 ┆ i64 │
╞═════╡
│ 2 ┆ 0 │
│ 6 ┆ 1 │
│ 7 ┆ 2 │
│ 3 ┆ 0 │
│ 9 ┆ 1 │
│ 2 ┆ 2 │
└─────┴

Installed versions

--------Version info---------
Polars:               1.1.0
Index type:           UInt32
Platform:             Linux-6.6.34-1-lts-x86_64-with-glibc2.39
Python:               3.12.4 (main, Jun  7 2024, 06:33:07) [GCC 14.1.1 20240522]

----Optional dependencies----
adbc_driver_manager:  1.1.0
cloudpickle:          <not installed>
connectorx:           0.3.3
deltalake:            <not installed>
fastexcel:            <not installed>
fsspec:               2024.6.1
gevent:               24.2.1
great_tables:         0.9.0
hvplot:               0.10.0
matplotlib:           3.9.0
nest_asyncio:         1.6.0
numpy:                1.26.4
openpyxl:             <not installed>
pandas:               2.2.2
pyarrow:              15.0.0
pydantic:             2.7.3
pyiceberg:            <not installed>
sqlalchemy:           2.0.31
torch:                <not installed>
xlsx2csv:             <not installed>
xlsxwriter:           <not installed>