feat: Add {DataFrame, LazyFrame}.cast

docs/api-reference/dataframe.md

@@ -6,6 +6,7 @@
       members:
         - __arrow_c_stream__
         - __getitem__
+        - cast
         - clone
         - collect_schema
         - columns

docs/api-reference/lazyframe.md

@@ -4,6 +4,7 @@
     handler: python
     options:
       members:
+        - cast
         - collect
         - collect_schema
         - columns

packages/test-plugin/src/test_plugin/dataframe.py

@@ -57,6 +57,7 @@ def _with_version(self, version: Version) -> Self:
 
     # Functions
     aggregate = not_implemented()
+    cast = not_implemented()
     collect = not_implemented()
     collect_schema = not_implemented()
     drop = not_implemented()

src/narwhals/_arrow/dataframe.py

@@ -708,6 +708,17 @@ def rename(self, mapping: Mapping[str, str]) -> Self:
             names = [mapping.get(c, c) for c in self.columns]
         return self._with_native(self.native.rename_columns(names))
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        native = self.native
+        # `with_type` keeps each field's name, nullability and metadata.
+        target = pa.schema(
+            field.with_type(narwhals_to_native_dtype(dtypes[field.name], self._version))
+            if field.name in dtypes
+            else field
+            for field in native.schema
+        )
+        return self._with_native(native.cast(target), validate_column_names=False)
+
     def write_parquet(self, file: str | Path | BytesIO) -> None:
         import pyarrow.parquet as pp
 

src/narwhals/_compliant/dataframe.py

@@ -151,6 +151,7 @@ def join_asof(
         suffix: str,
     ) -> Self: ...
     def rename(self, mapping: Mapping[str, str]) -> Self: ...
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self: ...
     def select(self, *exprs: CompliantExprT_contra) -> Self: ...
     def simple_select(self, *column_names: str) -> Self:
         """`select` where all args are column names."""

src/narwhals/_dask/dataframe.py

@@ -4,7 +4,7 @@
 
 import dask.dataframe as dd
 
-from narwhals._dask.utils import add_row_index, evaluate_exprs
+from narwhals._dask.utils import add_row_index, evaluate_exprs, narwhals_to_native_dtype
 from narwhals._pandas_like.utils import native_to_narwhals_dtype, select_columns_by_name
 from narwhals._typing_compat import assert_never
 from narwhals._utils import (
@@ -39,7 +39,12 @@
     from narwhals.dataframe import LazyFrame
     from narwhals.dtypes import DType
     from narwhals.exceptions import ColumnNotFoundError
-    from narwhals.typing import AsofJoinStrategy, JoinStrategy, UniqueKeepStrategy
+    from narwhals.typing import (
+        AsofJoinStrategy,
+        IntoDType,
+        JoinStrategy,
+        UniqueKeepStrategy,
+    )
 
 Incomplete: TypeAlias = "Any"
 """Using `_pandas_like` utils with `_dask`.
@@ -240,6 +245,13 @@ def with_row_index(self, name: str, order_by: Sequence[str] | None) -> Self:
     def rename(self, mapping: Mapping[str, str]) -> Self:
         return self._with_native(self.native.rename(columns=mapping))
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        native_dtypes = {
+            name: narwhals_to_native_dtype(dtype, self._version)
+            for name, dtype in dtypes.items()
+        }
+        return self._with_native(self.native.astype(native_dtypes))
+
     def head(self, n: int) -> Self:
         return self._with_native(self.native.head(n=n, compute=False, npartitions=-1))
 

src/narwhals/_ibis/dataframe.py

@@ -8,7 +8,12 @@
 import ibis.expr.types as ir
 
 from narwhals._ibis.expr import IbisExpr
-from narwhals._ibis.utils import evaluate_exprs, lit, native_to_narwhals_dtype
+from narwhals._ibis.utils import (
+    evaluate_exprs,
+    lit,
+    narwhals_to_native_dtype,
+    native_to_narwhals_dtype,
+)
 from narwhals._sql.dataframe import SQLLazyFrame
 from narwhals._utils import (
     Implementation,
@@ -41,7 +46,12 @@
     from narwhals.dataframe import LazyFrame
     from narwhals.dtypes import DType
     from narwhals.stable.v1 import DataFrame as DataFrameV1
-    from narwhals.typing import AsofJoinStrategy, JoinStrategy, UniqueKeepStrategy
+    from narwhals.typing import (
+        AsofJoinStrategy,
+        IntoDType,
+        JoinStrategy,
+        UniqueKeepStrategy,
+    )
 
     JoinPredicates: TypeAlias = "Sequence[ir.BooleanColumn] | Sequence[str]"
 
@@ -181,6 +191,14 @@ def with_columns(self, *exprs: IbisExpr) -> Self:
         new_columns_map = dict(evaluate_exprs(self, *exprs))
         return self._with_native(self.native.mutate(**new_columns_map))
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        # Ibis tables cast natively; a partial mapping leaves other columns untouched.
+        native_dtypes = {
+            name: narwhals_to_native_dtype(dtype, self._version)
+            for name, dtype in dtypes.items()
+        }
+        return self._with_native(self.native.cast(native_dtypes))
+
     def filter(self, predicate: IbisExpr) -> Self:
         # `[0]` is safe as the predicate's expression only returns a single column
         mask = cast("ir.BooleanValue", predicate(self)[0])

src/narwhals/_pandas_like/dataframe.py

@@ -521,6 +521,20 @@ def rename(self, mapping: Mapping[str, str]) -> Self:
             rename(self.native, columns=mapping, implementation=self._implementation)
         )
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        native = self.native
+        pd_dtypes = native.dtypes
+        to_cast = {
+            name: narwhals_to_native_dtype(
+                dtype,
+                dtype_backend=get_dtype_backend(pd_dtypes[name], self._implementation),
+                implementation=self._implementation,
+                version=self._version,
+            )
+            for name, dtype in dtypes.items()
+        }
+        return self._with_native(native.astype(to_cast), validate_column_names=False)
+
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
         to_drop = parse_columns_to_drop(self, columns, strict=strict)
         return self._with_native(

src/narwhals/_polars/dataframe.py

@@ -167,6 +167,15 @@ def __native_namespace__(self) -> ModuleType:
     def _with_native(self, df: NativePolarsFrame) -> Self:
         return self.__class__(df, version=self._version)
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        native_dtypes = {
+            name: narwhals_to_native_dtype(dtype, self._version)
+            for name, dtype in dtypes.items()
+        }
+        return self._with_native(
+            self.native.cast(native_dtypes)  # type: ignore[arg-type]
+        )
+
     def _with_version(self, version: Version) -> Self:
         return self.__class__(self.native, version=version)
 

src/narwhals/_sql/dataframe.py

@@ -13,14 +13,15 @@
 from narwhals.exceptions import MultiOutputExpressionError
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import Mapping, Sequence
     from typing import TypeAlias
 
     from typing_extensions import Self
 
     from narwhals._compliant.window import WindowInputs
     from narwhals._sql.expr import SQLExpr
     from narwhals.exceptions import ColumnNotFoundError
+    from narwhals.typing import IntoDType
 
     Incomplete: TypeAlias = Any
 
@@ -64,3 +65,15 @@ def filter(self, predicate: CompliantExprT_contra) -> Self:
             filtered = lf_with_tmp._filter(ns.col(tmp_col))
             return filtered.drop([tmp_col], strict=False)
         return self._filter(predicate)
+
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        if not dtypes:
+            # An empty mapping is a no-op, and it has to be answered as one
+            # here rather than passed through: `with_columns()` with nothing
+            # to add asks the engine to plan a projection over no columns, and
+            # PySpark Connect asserts against that in `plan.py`.
+            return self
+        ns = self.__narwhals_namespace__()
+        return self.with_columns(
+            *(ns.col(name).cast(dtype) for name, dtype in dtypes.items())
+        )

src/narwhals/dataframe.py

@@ -247,6 +247,11 @@ def select(
     def rename(self, mapping: dict[str, str]) -> Self:
         return self._with_compliant(self._compliant_frame.rename(mapping))
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        if error := self._check_columns_exist(list(dtypes)):
+            raise error
+        return self._with_compliant(self._compliant_frame.cast(dtypes))
+
     def head(self, n: int) -> Self:
         return self._with_compliant(self._compliant_frame.head(n))
 
@@ -1571,6 +1576,27 @@ def rename(self, mapping: dict[str, str]) -> Self:
         """
         return super().rename(mapping)
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        """Cast columns to the given dtypes.
+
+        Arguments:
+            dtypes: Mapping from column name to the dtype to cast it to. Columns not
+                in the mapping are left unchanged.
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]})
+            >>> nw.from_native(df_native).cast({"bar": nw.Int32}).to_native()
+            pyarrow.Table
+            foo: int64
+            bar: int32
+            ----
+            foo: [[1,2]]
+            bar: [[6,7]]
+        """
+        return super().cast(dtypes)
+
     def head(self, n: int = 5) -> Self:
         """Get the first `n` rows.
 
@@ -2420,6 +2446,10 @@ def _validate_metadata(self, metadata: ExprMetadata) -> None:
             )
             raise InvalidOperationError(msg)
 
+    def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None:
+        # `collect_schema` avoids the warning `self.columns` raises on a LazyFrame.
+        return check_columns_exist(subset, available=self.collect_schema().names())
+
     def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) -> None:
         self._level = level
         self._compliant_frame: CompliantLazyFrame[Any, LazyFrameT, Self]
@@ -2786,6 +2816,32 @@ def rename(self, mapping: dict[str, str]) -> Self:
         """
         return super().rename(mapping)
 
+    def cast(self, dtypes: Mapping[str, IntoDType]) -> Self:
+        r"""Cast columns to the given dtypes.
+
+        Arguments:
+            dtypes: Mapping from column name to the dtype to cast it to. Columns not
+                in the mapping are left unchanged.
+
+        Examples:
+            >>> import duckdb
+            >>> import narwhals as nw
+            >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 4.5), (3, 2.) df(a, b)")
+            >>> nw.from_native(lf_native).cast({"a": nw.Int64, "b": nw.Float64})
+            ┌──────────────────┐
+            |Narwhals LazyFrame|
+            |------------------|
+            |┌───────┬────────┐|
+            |│   a   │   b    │|
+            |│ int64 │ double │|
+            |├───────┼────────┤|
+            |│     1 │    4.5 │|
+            |│     3 │    2.0 │|
+            |└───────┴────────┘|
+            └──────────────────┘
+        """
+        return super().cast(dtypes)
+
     def head(self, n: int = 5) -> Self:
         r"""Get `n` rows.
 

tests/frame/cast_test.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import pytest
+
+import narwhals as nw
+from tests.utils import Constructor, assert_equal_data
+
+
+def test_cast(constructor: Constructor) -> None:
+    data = {"a": [1, 2, 3], "b": [4.0, 5.0, 6.0], "c": ["x", "y", "z"]}
+    df = nw.from_native(constructor(data))
+    original = df.collect_schema()
+
+    result = df.cast({"a": nw.Float64, "b": nw.Int32})
+    schema = result.collect_schema()
+    assert schema["a"] == nw.Float64
+    assert schema["b"] == nw.Int32
+    # A column that is not in the mapping keeps its dtype.
+    assert schema["c"] == original["c"]
+    assert_equal_data(
+        result, {"a": [1.0, 2.0, 3.0], "b": [4, 5, 6], "c": ["x", "y", "z"]}
+    )
+
+
+def test_cast_empty_mapping(constructor: Constructor) -> None:
+    data = {"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]}
+    df = nw.from_native(constructor(data))
+    result = df.cast({})
+    assert result.collect_schema() == df.collect_schema()
+    assert_equal_data(result, data)
+
+
+def test_cast_nonexistent_column(constructor: Constructor) -> None:
+    df = nw.from_native(constructor({"a": [1, 2, 3]}))
+    with pytest.raises(nw.exceptions.ColumnNotFoundError):
+        df.cast({"b": nw.Int64})