Type casting issue when select is called with complex expression: pyo3_runtime.PanicException: called Option::unwrap() on a None value

[bike-picket]

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

df = pl.from_dicts(
    {
        "a": [1, 2],
        "b": [99, 99],
    }
).lazy()
df = df.with_row_index().select(
    pl.when(pl.col("b") < 10).then(0.1 * pl.col("b")).when(pl.col("b") < 100).then(0.2 * pl.col("b"))
)
res = df.collect().to_dicts()  # pyo3_runtime.PanicException: called `Option::unwrap()` on a `None` value

Log output

thread '<unnamed>' panicked at crates/polars-lazy/src/physical_plan/expressions/column.rs:127:14:
called `Option::unwrap()` on a `None` value
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

Issue description

Interestingly, the issue is not reproduced if:
a) with_row_index call is removed;
b) or b column is initialized as float ("b": [99.0, 99.0]).

Expected behavior

The code should be executed successfully and return:

[{'literal': 19.8}, {'literal': 19.8}]

Installed versions

None

[cmdlineluser]

Can reproduce. Seems to be a CSE bug.

import polars as pl

df = pl.LazyFrame({
    "a": [1, 2],
    "b": [99, 99]
}).with_row_index()

df.select(
    pl.when(pl.col("b") < 10)
      .then(0.1 * pl.col("b"))
      .when(pl.col("b") < 100)
      .then(0.2 * pl.col("b"))
).collect(comm_subexpr_elim=False)

# shape: (2, 1)
# ┌─────────┐
# │ literal │
# │ ---     │
# │ f64     │
# ╞═════════╡
# │ 19.8    │
# │ 19.8    │
# └─────────┘

[deanm0000]

It seems to be getting "confused" by the with_row_index so even though the select should only return 1 column, you can see in the explain that it wants to return the result of the when/then and just col('b') by itself but with the alias of the cse.

df = pl.LazyFrame({
    "a": [1, 2],
    "b": [99, 99]
})
print(
df.lazy().with_row_index().select(
    pl.when(pl.col("b") < 10).then(0.1 * pl.col("b")).when(pl.col("b") < 100).then(0.2 * pl.col("b"))
).explain(optimized=True))
 SELECT [.when([(col("b")) < (10)]).then([(0.1) * (col("__POLARS_CSER_704771700852997304"))]).otherwise(.when([(col("b")) < (100)]).then([(0.2) * (col("__POLARS_CSER_704771700852997304"))]).otherwise(null.cast(Float64))).alias("literal"), col("b").cast(Float64).alias("__POLARS_CSER_704771700852997304")] FROM
  WITH ROW INDEX
    DF ["a", "b"]; PROJECT 1/2 COLUMNS; SELECTION: "None"

Given that, it's not surprising that if we turn off projection_pushdown then it will also work

print(
df.lazy().with_row_index().select(
    pl.when(pl.col("b") < 10).then(0.1 * pl.col("b")).when(pl.col("b") < 100).then(0.2 * pl.col("b"))
).collect(projection_pushdown=False))
shape: (2, 1)
┌─────────┐
│ literal │
│ ---     │
│ f64     │
╞═════════╡
│ 19.8    │
│ 19.8    │
└─────────┘

If we replace the with_row_index with with_columns(index=pl.int_range(0,pl.len())) then it will correctly drop that from the query plan while still using CSE and it works

print(
df.lazy().with_columns(index=pl.int_range(0,pl.len())).select(
    pl.when(pl.col("b") < 10).then(0.1 * pl.col("b")).when(pl.col("b") < 100).then(0.2 * pl.col("b"))
).explain(optimized=True))
 SELECT [.when([(col("b")) < (10)]).then([(0.1) * (col("__POLARS_CSER_704771700852997304"))]).otherwise(.when([(col("b")) < (100)]).then([(0.2) * (col("__POLARS_CSER_704771700852997304"))]).otherwise(null.cast(Float64))).alias("literal"), col("b").cast(Float64).alias("__POLARS_CSER_704771700852997304")] FROM
  DF ["a", "b"]; PROJECT 1/2 COLUMNS; SELECTION: "None"

I marked this as high priority because it's so interesting how many things have to be happening for this bug to come out and how many unique workarounds there are that all seem so unrelated