facebook / facebook/pyrefly

Pyrefly loses Polars DataFrame schemas for positional `with_columns` and keyword `select`

Open
#4,565 4 comments 0 reactions 1 assignee Claimed by @Vishwaspatel2401 View on GitHub
dataframe
Dominant language
Rust
Stars
7k
Forks
516
PR merge metrics
No merged PRs in 30d

Description

### Describe the Bug

`select` and `with_columns` support opposite expression forms. `select` preserves the schema for positional expressions but not keyword expressions, while `with_columns` does the reverse.

```python
from typing import reveal_type

import polars as pl

df = pl.DataFrame({"units": [12]})

reveal_type(df.select(pl.col("units").alias("qty"))) # DataFrame[qty: Int64]
reveal_type(df.with_columns(pl.col("units").alias("qty"))) # DataFrame
reveal_type(df.select(qty=pl.col("units"))) # DataFrame
reveal_type(df.with_columns(qty=pl.col("units"))) # DataFrame[units: Int64, qty: Int64]
```

I would expect the equivalent forms to produce the same schemas:

```text
select(...) -> DataFrame[qty: Int64]
with_columns(...) -> DataFrame[units: Int64, qty: Int64]
```

The common positional form of `with_columns` therefore loses all downstream column checking.
This also affects conditional positional expressions:

```python
reveal_type(
df.with_columns(
pl.when(pl.col("units") > 10)
.then(pl.lit("high"))
.otherwise(pl.lit("low"))
.alias("bucket")
)
) # DataFrame
```

Reproduced on `pyrefly 1.3.0-dev.1` with Polars 1.35.2.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.