[MarcoGorelli]

Summary

Fixes #4976. The general idea is to use DataFrameKind::Narwhals to piggy-back off of the existing Polars support, and only track computed dtypes in cases where Narwhals guarantees them (which is fewer than when Polars does)

Demo

What works well

Let's start with

from typing import Annotated, reveal_type

import narwhals as nw
import polars as pl


class Sales:
    region: nw.String
    units: nw.Int64
    price: nw.Float64


class Regions:
    region: nw.String
    manager: nw.String


def entry(df: Annotated[nw.DataFrame[pl.DataFrame], Sales]) -> None:
    reveal_type(df)
    df["not_declared"]

def partial(df: Annotated[nw.DataFrame, Sales, ...]) -> None:
    reveal_type(df)
    df["not_declared"]

Note that Narwhals dataframes are themselves generic over the native library. We can either write nw.DataFrame[pl.DataFrame], or just nw.DataFrame and leave the native library unknown

Pyrefly shows the expected types:

• narwhals.dataframe.DataFrame[polars.dataframe.frame.DataFrame][region: String, units: Int64, price: Float64]
• DataFrame[Unknown][region: String, units: Int64, price: Float64, ...]

It also correctly finds the error in entry:

ERROR Column `not_declared` is not in the DataFrame schema [unknown-column]
  --> t.py:22:8
   |
22 |     df["not_declared"]

but is fine with it in partial ­ 👍

Series types are detected. Adding reveal_type(df['region']) to entry shows

revealed type: Series[Any][String]

Projections are fine / work as expected in Polars

def projections(df: Annotated[nw.DataFrame, Sales]) -> None:
    reveal_type(df.select("region", "units"))  # [region: String, units: Int64]
    reveal_type(df.select(nw.col("price")))  # [price: Float64]
    reveal_type(df[["region", "price"]])  # [region: String, price: Float64]
    reveal_type(df.drop("price"))  # [region: String, units: Int64]
    reveal_type(df.rename({"units": "qty"}))  # [region: String, qty: Int64, price: Float64]
    reveal_type(df.with_columns(revenue=nw.col("units") * nw.col("price")))  # [..., revenue: Float64]
    df.select("revenue")  # ERROR unknown-column
    df.drop("revenue")  # ERROR unknown-column
    df.select("units", "units")  # ERROR duplicate-column

A few more demos:

# ---------------------------------------------------------------------------
# 4. Row transforms keep the schema and check column names.
# ---------------------------------------------------------------------------
def rows(df: Annotated[nw.DataFrame, Sales]) -> None:
    reveal_type(df.filter(nw.col("units") > 0))  # schema preserved
    reveal_type(df.filter(region="EU"))  # schema preserved
    reveal_type(df.sort("price", descending=True))  # schema preserved
    reveal_type(df.unique(subset=["region"]))  # schema preserved
    reveal_type(df.drop_nulls(subset="price"))  # schema preserved
    reveal_type(df.head(5))  # schema preserved
    df.sort("revenue")  # ERROR unknown-column
    df.unique(subset=["revenue"])  # ERROR unknown-column
    df.filter(revenue=1)  # ERROR unknown-column


# ---------------------------------------------------------------------------
# 5. group_by / agg
# ---------------------------------------------------------------------------
def aggregations(df: Annotated[nw.DataFrame, Sales]) -> None:
    reveal_type(df.group_by("region").agg(nw.col("units").sum()))  # [region: String, units: Unknown]
    reveal_type(df.group_by("region").agg(nw.col("price").mean().alias("avg")))  # [region: String, avg: Float64]
    df.group_by("revenue").agg(nw.len())  # ERROR unknown-column


# ---------------------------------------------------------------------------
# 6. Joins and concat.
# ---------------------------------------------------------------------------
def combine(
    sales: Annotated[nw.DataFrame, Sales],
    more: Annotated[nw.DataFrame, Sales],
    regions: Annotated[nw.DataFrame, Regions],
) -> None:
    reveal_type(sales.join(regions, on="region"))  # [region, units, price, manager: String]
    reveal_type(nw.concat([sales, more]))  # DataFrame[Any][region: String, units: Int64, price: Float64]


# ---------------------------------------------------------------------------
# 7. Lazy frames: lazy() / collect() round-trip, and annotated LazyFrames.
# ---------------------------------------------------------------------------
def lazy(df: Annotated[nw.DataFrame, Sales], lf: Annotated[nw.LazyFrame, Sales]) -> None:
    reveal_type(df.lazy())  # LazyFrame[Any][region: String, units: Int64, price: Float64]
    reveal_type(df.lazy().select("units").collect())  # DataFrame[Any][units: Int64]
    reveal_type(lf)  # LazyFrame[Unknown][region: String, units: Int64, price: Float64]
    reveal_type(lf.group_by("region").agg(nw.col("units").sum()))  # LazyFrame[Unknown][region: String, units: Unknown]
    lf.select("revenue")  # ERROR unknown-column

# ---------------------------------------------------------------------------
# 8. Schema contracts: return types are checked, and frame libraries don't mix.
# ---------------------------------------------------------------------------
class Summary:
    region: nw.String
    units: nw.Int64


def summarise(df: Annotated[nw.DataFrame, Sales]) -> Annotated[nw.DataFrame, Summary]:
    return df.select("region", "units")  # OK


def summarise_wrong(
    df: Annotated[nw.DataFrame, Sales],
) -> Annotated[nw.DataFrame, Summary]:
    return df.select("region")  # ERROR bad-return


def takes_narwhals(df: Annotated[nw.DataFrame, Summary]) -> None: ...


def mixing(pl_df: Annotated[pl.DataFrame, Summary]) -> None:
    takes_narwhals(pl_df)  # ERROR bad-argument-type

Stable namespaces (narwhals.stable.v1 / narwhals.stable.v2) are tracked the same way

What's left out

• Nothing is inferred in nw.from_dict / nw.from_native. Initially at least, I think this is fine. If there's interest, it could be addressed as a follow-up
• It's technically possible to annotate a Narwhals dataframe using Polars dtypes. I don't think any validation is done the types currently, so this may be best left separate from this PR?
• There's fewer rules about types resulting from operations (e.g. nw.len()) than there are in Polars, as the result in Narwhals may vary depending on the backend. I think this is OK, it's easier to start more relaxed and become stricter later than to do the other way round

Test Plan

tests are included

[meta-codesync]

This pull request has been imported. If you are a Meta employee, you can view this in D120721406. (Because this pull request was imported automatically, there will not be any future comments.)

[stroxler]

This basically looks good to me.

Back to you for two minor things:

• A rebase would help me merge, there's now a conflict with trunk (probably very minor)
• Should we have test stubs for the v2 API? I don't know much about this so I trust you, but an agent review flagged that this might be worth asking about

[github-actions]

According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅

[MarcoGorelli]

thanks for your review! sure, done

[stroxler]

Review automatically exported from Phabricator review in Meta.