feat: Add {DataFrame, LazyFrame}.cast - #3815
Conversation
{DataFrame, LazyFrame}.cast
There was a problem hiding this comment.
Thanks for the contribution @mkzung - I left a few inline comments. On top of those, I wonder about the rationale regarding:
the SQL backends (duckdb, pyspark) cast per column through with_columns, and ibis inherits that path
as listed in the issue, pyspark and ibis have their own casting method and I think we should used them instead of choosing for a user. DuckDB doesn't seem to have it, SQLFrame I didn't check.
Also, there are a few conflicts to solve with the main branch before being able to merge
| for name, dtype in dtypes.items() | ||
| } | ||
| return self._with_native( | ||
| self.native.cast(native_dtypes) # pyright: ignore[reportArgumentType] |
There was a problem hiding this comment.
Is # pyright: ignore[reportArgumentType] an upstream polars issue?
|
Rebased, conflicts resolved. On the comments:
Also moved the LazyFrame column check to |
|
Thanks for adjusting @mkzung
Preserving mutability and metadata is a good point. I wonder if we should do it in any cast operation.
What about Can you also please run all the checks locally as mentioned in the contributing guide? |
|
Benchmark script and numbers (pyarrow 23, 500k rows x 8 int columns, 50 runs after warm-up): import time
import numpy as np
import pyarrow as pa
import narwhals as nw
rng = np.random.default_rng(0)
tbl = pa.table({f"c{i}": rng.integers(0, 100, 500_000) for i in range(8)})
df = nw.from_native(tbl, eager_only=True)
targets = {f"c{i}": nw.Float64 for i in range(8)}
def schema_cast():
return df.cast(targets)
def with_columns_cast():
return df.with_columns(*(nw.col(n).cast(d) for n, d in targets.items()))
for fn in (schema_cast, with_columns_cast):
fn()
start = time.perf_counter()
for _ in range(50):
fn()
print(f"{fn.__name__}: {(time.perf_counter() - start) / 50 * 1000:.2f} ms")So no meaningful speed difference; the reason to keep the schema path is only the On On doing the field-preserving cast everywhere: for the eager per-column Checks: prek, pyright/mypy/pyrefly and the full-coverage suite were run locally before the |
Cast columns to given dtypes from a {name: dtype} mapping, leaving columns not in
the mapping unchanged. Uses the native frame cast on the eager backends (pyarrow
Table.cast, pandas astype, polars cast) and dask astype; the SQL backends (duckdb,
pyspark, ibis) cast per column through with_columns.
Closes narwhals-dev#3402
…e lookup - ibis: override cast to use the native ir.Table.cast (partial mapping leaves other columns untouched) instead of the shared per-column with_columns path - override _check_columns_exist on LazyFrame to resolve names via collect_schema().names(); this avoids the column-access PerformanceWarning that self.columns raises on a polars LazyFrame, while the eager DataFrame keeps the cheap self.columns path - pandas-like: read the column dtype from native.dtypes[name] instead of slicing the frame with native[name] - polars: suppress the cast mapping arg-type stub gap with the backend's # type: ignore[arg-type] convention
PySpark Connect asserts in plan.py when asked to build a projection over no
columns, so cast({}) failed there while every other backend treated it as the
no-op it is. The SQL backend returns self for an empty mapping now, which is
also one fewer plan node on the backends that did tolerate it.
Adding cast to the compliant frame protocol made DictLazyFrame abstract, so pyright refused to instantiate it in namespace.py. Declared alongside the other methods the plugin does not implement.
Description
Adds
DataFrame.castandLazyFrame.cast.Takes a
{column: dtype}mapping and casts those columns, leaving the rest unchanged. It uses each backend's native frame cast where there is one (pyarrowTable.cast,pandas/daskastype,polarscast); the SQL backends (duckdb,pyspark) cast per column throughwith_columns, andibisinherits that path. Casting a column that is not in the frame raisesColumnNotFoundError, the same asselectanddrop.I kept it to the mapping form, which is what was proposed in the issue. Polars also takes a single dtype to cast every column, so I can add that too if you'd like it for parity.
What type of PR is this? (check all applicable)
Related issues
{DataFrame, LazyFrame}.cast(...)#3402AI assistance
Checklist