filtering a concatenated LazyFrame raises incorrect ColumnNotFoundError exception

[erinov1]

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

N = 2

df = pl.LazyFrame({"x": "a", "y": 1})
df = pl.concat(
    df.with_columns(pl.col("y").add(n)) for n in range(N)
)

df.filter(pl.col("x").eq("a")).collect()

Log output

found multiple sources; run comm_subplan_elim
UNION: `parallel=false` union is run sequentially
CACHE SET: cache id: 0
---------------------------------------------------------------------------
ColumnNotFoundError                       Traceback (most recent call last)
Cell In[8], line 8
      3 df = pl.LazyFrame({"x": "a", "y": 1})
      4 df = pl.concat(
      5     df.with_columns(pl.col("y").add(n)) for n in range(N)
      6 )
----> 8 df.filter(pl.col("x").eq("a")).collect()

File ~/Library/Caches/pypoetry/virtualenvs/ijk-Wd_A0uOh-py3.11/lib/python3.11/site-packages/polars/lazyframe/frame.py:1810, in LazyFrame.collect(self, type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, no_optimization, streaming, background, _eager)
   1807 if background:
   1808     return InProcessQuery(ldf.collect_concurrently())
-> 1810 return wrap_df(ldf.collect())

ColumnNotFoundError: y

Issue description

The query above fails for N > 1. Exchanging the order of the filter and final collect also works correctly, or if working eagerly throughout.

Expected behavior

The query should run and return

shape: (2, 2)
┌─────┬─────┐
│ x   ┆ y   │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ a   ┆ 1   │
│ a   ┆ 2   │
└─────┴─────┘

like in eager mode.

Installed versions

--------Version info---------
Polars:               0.20.23
Index type:           UInt32
Platform:             macOS-14.4.1-arm64-arm-64bit
Python:               3.11.1 (main, Feb  8 2023, 16:32:14) [Clang 14.0.0 (clang-1400.0.29.202)]

----Optional dependencies----
adbc_driver_manager:  <not installed>
cloudpickle:          <not installed>
connectorx:           <not installed>
deltalake:            <not installed>
fastexcel:            <not installed>
fsspec:               2024.3.1
gevent:               <not installed>
hvplot:               <not installed>
matplotlib:           3.7.5
nest_asyncio:         1.6.0
numpy:                1.24.4
openpyxl:             <not installed>
pandas:               1.5.3
pyarrow:              12.0.1
pydantic:             1.10.5
pyiceberg:            <not installed>
pyxlsb:               <not installed>
sqlalchemy:           2.0.29
xlsx2csv:             <not installed>
xlsxwriter:           <not installed>

[cmdlineluser]

Can reproduce.

Seems to be one of those optimizer issues:

>>> df.filter(pl.col("x").eq("a")).collect(comm_subplan_elim=False)
shape: (2, 2)
┌─────┬─────┐
│ x   ┆ y   │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ a   ┆ 1   │
│ a   ┆ 2   │
└─────┴─────┘

>>> df.filter(pl.col("x").eq("a")).collect(predicate_pushdown=False)
shape: (2, 2)
┌─────┬─────┐
│ x   ┆ y   │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ a   ┆ 1   │
│ a   ┆ 2   │
└─────┴─────┘

[knl]

@ritchie46 What would be a workaround for this, for people that can't upgrade to the version that fixes the issue?

EDIT: I see that swapping collect and filter would do it.

[cmdlineluser]

@knl You can try disabling specific optimizations to check for potential causes e.g. .collect(comm_subplan_elim=False)

The example query in this issue ran for me under 2 situations: #15980 (comment)