diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md index b21ce3cea9..9a11e87cb9 100644 --- a/docs/api-reference/dataframe.md +++ b/docs/api-reference/dataframe.md @@ -6,6 +6,7 @@ members: - __arrow_c_stream__ - __getitem__ + - cast - clone - collect_schema - columns diff --git a/docs/api-reference/lazyframe.md b/docs/api-reference/lazyframe.md index f9fafb7145..d4749dc962 100644 --- a/docs/api-reference/lazyframe.md +++ b/docs/api-reference/lazyframe.md @@ -4,6 +4,7 @@ handler: python options: members: + - cast - collect - collect_schema - columns diff --git a/packages/test-plugin/src/test_plugin/dataframe.py b/packages/test-plugin/src/test_plugin/dataframe.py index 56bc928cf7..0af402c0e1 100644 --- a/packages/test-plugin/src/test_plugin/dataframe.py +++ b/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() diff --git a/src/narwhals/_arrow/dataframe.py b/src/narwhals/_arrow/dataframe.py index 3a6e5aa4ae..3255a10459 100644 --- a/src/narwhals/_arrow/dataframe.py +++ b/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 diff --git a/src/narwhals/_compliant/dataframe.py b/src/narwhals/_compliant/dataframe.py index d4c1109e1b..2574cf9df5 100644 --- a/src/narwhals/_compliant/dataframe.py +++ b/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.""" diff --git a/src/narwhals/_dask/dataframe.py b/src/narwhals/_dask/dataframe.py index 809e94372a..0c9861d5f8 100644 --- a/src/narwhals/_dask/dataframe.py +++ b/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)) diff --git a/src/narwhals/_ibis/dataframe.py b/src/narwhals/_ibis/dataframe.py index 97587c5dbf..752e21ab45 100644 --- a/src/narwhals/_ibis/dataframe.py +++ b/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]) diff --git a/src/narwhals/_pandas_like/dataframe.py b/src/narwhals/_pandas_like/dataframe.py index 52f5a960ae..5923e805b4 100644 --- a/src/narwhals/_pandas_like/dataframe.py +++ b/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( diff --git a/src/narwhals/_polars/dataframe.py b/src/narwhals/_polars/dataframe.py index 7b96aceaa2..8fb7d6aec1 100644 --- a/src/narwhals/_polars/dataframe.py +++ b/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) diff --git a/src/narwhals/_sql/dataframe.py b/src/narwhals/_sql/dataframe.py index 7a812875c6..2afe9b3a63 100644 --- a/src/narwhals/_sql/dataframe.py +++ b/src/narwhals/_sql/dataframe.py @@ -13,7 +13,7 @@ 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 @@ -21,6 +21,7 @@ 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()) + ) diff --git a/src/narwhals/dataframe.py b/src/narwhals/dataframe.py index e18bb409f9..090f387a28 100644 --- a/src/narwhals/dataframe.py +++ b/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. diff --git a/tests/frame/cast_test.py b/tests/frame/cast_test.py new file mode 100644 index 0000000000..f39135e7fc --- /dev/null +++ b/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})