Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 34 additions & 2 deletions src/narwhals/_arrow/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,22 @@ def maybe_extract_py_scalar(value: Any, return_py_scalar: bool) -> Any: # noqa:
return value


def _cast_scalar_non_strict(value: ScalarAny, data_type: pa.DataType) -> ScalarAny:
try:
return pc.cast(value, data_type)
except pa.ArrowInvalid:
return pa.scalar(None, type=data_type)


def _cast_chunk_non_strict(chunk: ArrayAny, data_type: pa.DataType) -> ArrayAny:
try:
return pc.cast(chunk, data_type)
except pa.ArrowInvalid:
return pa.array(
[_cast_scalar_non_strict(value, data_type) for value in chunk], type=data_type
)


class ArrowSeries(EagerSeries["ChunkedArrayAny"]):
_implementation = Implementation.PYARROW

Expand Down Expand Up @@ -546,9 +562,25 @@ def is_null(self) -> Self:
def is_nan(self) -> Self:
return self._with_native(pc.is_nan(self.native), preserve_broadcast=True)

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
data_type = narwhals_to_native_dtype(dtype, self._version)
return self._with_native(pc.cast(self.native, data_type), preserve_broadcast=True)
if strict:
native = pc.cast(self.native, data_type)
else:
dtypes = self._version.dtypes
base_type = dtype.base_type()
is_fallible_cast = base_type.is_numeric() or issubclass(
base_type, (dtypes.Date, dtypes.Datetime, dtypes.Duration)
)
if not is_fallible_cast:
native = pc.cast(self.native, data_type)
else:
chunks = [
_cast_chunk_non_strict(chunk, data_type)
for chunk in self.native.chunks
]
native = pa.chunked_array(chunks, type=data_type)
return self._with_native(native, preserve_broadcast=True)

def null_count(self, *, _return_py_scalar: bool = True) -> int:
return maybe_extract_py_scalar(self.native.null_count, _return_py_scalar)
Expand Down
2 changes: 1 addition & 1 deletion src/narwhals/_compliant/column.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ def __narwhals_namespace__(self) -> CompliantNamespace[Any, Any]: ...

def abs(self) -> Self: ...
def alias(self, name: str) -> Self: ...
def cast(self, dtype: IntoDType) -> Self: ...
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self: ...
def clip(self, lower_bound: Self, upper_bound: Self) -> Self: ...
def clip_lower(self, lower_bound: Self) -> Self: ...
def clip_upper(self, upper_bound: Self) -> Self: ...
Expand Down
4 changes: 2 additions & 2 deletions src/narwhals/_compliant/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -449,8 +449,8 @@ def func(df: EagerDataFrameT) -> list[EagerSeriesT]:
version=self._version,
)

def cast(self, dtype: IntoDType) -> Self:
return self._reuse_series("cast", dtype=dtype)
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
return self._reuse_series("cast", dtype=dtype, strict=strict)

def _with_binary(self, operator: str, other: Self, /) -> Self:
return self._reuse_series(operator, other=other)
Expand Down
6 changes: 5 additions & 1 deletion src/narwhals/_dask/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -618,7 +618,11 @@ def func(df: DaskLazyFrame) -> Sequence[dx.Series]:
version=self._version,
)

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
if not strict:
msg = "`cast(..., strict=False)` is not yet implemented for the dask backend."
raise NotImplementedError(msg)

def func(expr: dx.Series) -> dx.Series:
native_dtype = narwhals_to_native_dtype(dtype, self._version)
return expr.astype(native_dtype)
Expand Down
19 changes: 16 additions & 3 deletions src/narwhals/_duckdb/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,12 @@ def _any_value(self, expr: Expression, *, ignore_nulls: bool) -> Expression:
else self._function("first", expr)
)

@staticmethod
def _cast_non_strict(expr: Expression, native_dtype: Any) -> Expression:
# DuckDB's relational API has `Expression.cast`, but does not expose a
# TRY_CAST expression constructor.
return sql_expression(f"TRY_CAST({expr} AS {native_dtype})")

def __narwhals_namespace__(self) -> DuckDBNamespace: # pragma: no cover
from narwhals._duckdb.namespace import DuckDBNamespace

Expand Down Expand Up @@ -269,16 +275,23 @@ def _fill_constant(expr: Expression, value: Any) -> Expression:
assert value is not None # noqa: S101
return self._with_elementwise(_fill_constant, expression_args={"value": value})

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
def _cast(expr: Expression, native_dtype: Any) -> Expression:
if strict:
return expr.cast(native_dtype)
return self._cast_non_strict(expr, native_dtype)

def func(df: DuckDBLazyFrame) -> list[Expression]:
tz = DeferredTimeZone(df.native)
native_dtype = narwhals_to_native_dtype(dtype, self._version, tz)
return [expr.cast(native_dtype) for expr in self(df)]
return [_cast(expr, native_dtype) for expr in self(df)]

def window_f(df: DuckDBLazyFrame, inputs: DuckDBWindowInputs) -> list[Expression]:
tz = DeferredTimeZone(df.native)
native_dtype = narwhals_to_native_dtype(dtype, self._version, tz)
return [expr.cast(native_dtype) for expr in self.window_function(df, inputs)]
return [
_cast(expr, native_dtype) for expr in self.window_function(df, inputs)
]

return self.__class__(
func,
Expand Down
8 changes: 5 additions & 3 deletions src/narwhals/_ibis/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -285,11 +285,13 @@ def _fill_null(expr: ir.Value, value: ir.Scalar) -> ir.Value:
assert value is not None # noqa: S101
return self._with_callable(_fill_null, expression_args={"value": value})

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
def _func(expr: ir.Column) -> ir.Value:
native_dtype = narwhals_to_native_dtype(dtype, self._version)
# ibis `cast` overloads do not include DataType, only literals
return expr.cast(native_dtype) # type: ignore[unused-ignore]
# ibis `cast`/`try_cast` overloads do not include DataType, only literals
if strict:
return expr.cast(native_dtype) # type: ignore[unused-ignore]
return expr.try_cast(native_dtype) # type: ignore[unused-ignore]

return self._with_callable(_func)

Expand Down
31 changes: 31 additions & 0 deletions src/narwhals/_integer_bounds.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
"""Min/max representable values per integer dtype, shared across backends.

Used by non-strict (`strict=False`) casting to null-out (rather than silently wrap
or raise on) out-of-range values before handing off to a backend's native cast.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from narwhals._utils import Version

if TYPE_CHECKING:
from collections.abc import Mapping

from narwhals.dtypes import DType

__all__ = ["INTEGER_BOUNDS"]

_dtypes = Version.MAIN.dtypes

INTEGER_BOUNDS: Mapping[type[DType], tuple[int, int]] = {
_dtypes.Int8: (-128, 127),
_dtypes.Int16: (-32_768, 32_767),
_dtypes.Int32: (-2_147_483_648, 2_147_483_647),
_dtypes.Int64: (-(2**63), 2**63 - 1),
_dtypes.UInt8: (0, 255),
_dtypes.UInt16: (0, 65_535),
_dtypes.UInt32: (0, 2**32 - 1),
_dtypes.UInt64: (0, 2**64 - 1),
}
19 changes: 16 additions & 3 deletions src/narwhals/_pandas_like/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
align_and_extract_native,
binary_string_sum_fallback,
broadcast_series_to_index,
cast_non_strict,
get_dtype_backend,
import_array_module,
narwhals_to_native_dtype,
Expand Down Expand Up @@ -312,19 +313,31 @@ def scatter(

return None if in_place else self._with_native(series)

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
if self.dtype == dtype and self.native.dtype != "object":
# Avoid dealing with pandas' type-system if we can. Note that it's only
# safe to do this if we're not starting with object dtype, see tests/expr_and_series/cast_test.py::test_cast_object_pandas
# for an example of why.
return self._with_native(self.native, preserve_broadcast=True)
dtype_backend = get_dtype_backend(self.native.dtype, self._implementation)
pd_dtype = narwhals_to_native_dtype(
dtype,
dtype_backend=get_dtype_backend(self.native.dtype, self._implementation),
dtype_backend=dtype_backend,
implementation=self._implementation,
version=self._version,
)
return self._with_native(self.native.astype(pd_dtype), preserve_broadcast=True)
if strict:
native = self.native.astype(pd_dtype)
else:
native = cast_non_strict(
self.native,
dtype,
pd_dtype=pd_dtype,
dtype_backend=dtype_backend,
implementation=self._implementation,
version=self._version,
)
return self._with_native(native, preserve_broadcast=True)

def item(self, index: int | None = None) -> Any:
# cuDF doesn't have Series.item().
Expand Down
92 changes: 92 additions & 0 deletions src/narwhals/_pandas_like/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
US_PER_SECOND,
)
from narwhals._exceptions import issue_warning
from narwhals._integer_bounds import INTEGER_BOUNDS
from narwhals._utils import (
Implementation,
Version,
Expand Down Expand Up @@ -560,6 +561,97 @@ def narwhals_to_native_arrow_dtype(
raise NotImplementedError(msg)


def _cast_integer_non_strict(
native: pd.Series[Any],
dtype: IntoDType,
*,
pd_dtype: Any,
dtype_backend: DTypeBackend,
implementation: Implementation,
version: Version,
) -> pd.Series[Any]:
base_type = dtype.base_type()
source_is_numeric = pd.api.types.is_numeric_dtype(native.dtype)
numeric = pd.to_numeric(native, errors="coerce")
if (
not source_is_numeric
and pd.api.types.is_float_dtype(numeric.dtype)
and numeric.abs().gt(2**53).any()
):
msg = (
"Pandas converted integer-like input outside the exact Float64 range "
"while performing a non-strict cast. This conversion is not supported "
"because it may silently lose integer precision."
)
raise NotImplementedError(msg)

truncated = np.trunc(numeric)
if source_is_numeric:
# Match Pandas' native numeric-to-integer cast, which truncates fractions.
numeric = truncated
else:
# Pandas does not consider fractional text to be valid integer text.
numeric = numeric.where(numeric.isna() | numeric.eq(truncated))

lo, hi = INTEGER_BOUNDS[base_type]
numeric = numeric.where(numeric.isna() | numeric.between(lo, hi))
target: Any = pd_dtype
if (
dtype_backend is None
and implementation is not Implementation.CUDF
and numeric.isna().any()
):
# A plain NumPy integer dtype can't hold a null, so use Pandas'
# nullable representation when coercion introduces or preserves one.
target = narwhals_to_native_dtype(
dtype, "numpy_nullable", implementation, version
)
return numeric.astype(target)


def cast_non_strict(
native: pd.Series[Any],
dtype: IntoDType,
*,
pd_dtype: Any,
dtype_backend: DTypeBackend,
implementation: Implementation,
version: Version,
) -> pd.Series[Any]:
"""Non-strict (`strict=False`) counterpart to `.astype(pd_dtype)`-based casting.

Dispatches on the target dtype family and uses pandas' own vectorized "coerce"
tools, so that values which can't be cast become null instead of raising.
"""
base_type = dtype.base_type()
if base_type.is_integer():
return _cast_integer_non_strict(
native,
dtype,
pd_dtype=pd_dtype,
dtype_backend=dtype_backend,
implementation=implementation,
version=version,
)
if base_type.is_float():
numeric = pd.to_numeric(native, errors="coerce")
return numeric.astype(pd_dtype)
if issubclass(base_type, dtypes.Datetime):
converted = pd.to_datetime(native, errors="coerce")
# `to_datetime` parses at nanosecond resolution. Normalize before casting
# to an Arrow-backed coarser unit, which otherwise rejects lossy conversion.
converted = converted.dt.as_unit(dtype.time_unit)
return converted.astype(pd_dtype)
if issubclass(base_type, dtypes.Date):
return pd.to_datetime(native, errors="coerce").astype(pd_dtype)
if issubclass(base_type, dtypes.Duration):
return pd.to_timedelta(native, errors="coerce").astype(pd_dtype)
# Pandas has no coercing conversion API for the remaining target families.
# Delegate to its regular cast rather than inventing cross-backend semantics;
# unsupported conversions may therefore still raise with `strict=False`.
return native.astype(pd_dtype)


def int_dtype_mapper(dtype: Any) -> str:
if "pyarrow" in str(dtype):
return "Int64[pyarrow]"
Expand Down
4 changes: 2 additions & 2 deletions src/narwhals/_polars/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,9 +103,9 @@ def _renamed_min_periods(self, min_samples: int, /) -> dict[str, Any]:
name = "min_periods" if self._backend_version < (1, 21, 0) else "min_samples"
return {name: min_samples}

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
dtype_pl = narwhals_to_native_dtype(dtype, self._version)
return self._with_native(self.native.cast(dtype_pl))
return self._with_native(self.native.cast(dtype_pl, strict=strict))

def clip_lower(self, lower_bound: PolarsExpr) -> Self:
lower_native = extract_native(lower_bound)
Expand Down
4 changes: 2 additions & 2 deletions src/narwhals/_polars/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,9 +288,9 @@ def __getitem__(self, item: MultiIndexSelector[Self]) -> Any | Self:
return self._from_native_object(self.native.__getitem__(item.native))
return self._from_native_object(self.native.__getitem__(item))

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
dtype_pl = narwhals_to_native_dtype(dtype, self._version)
return self._with_native(self.native.cast(dtype_pl))
return self._with_native(self.native.cast(dtype_pl, strict=strict))

def clip(self, lower_bound: PolarsSeries, upper_bound: PolarsSeries) -> Self:
return self._with_native(
Expand Down
6 changes: 5 additions & 1 deletion src/narwhals/_spark_like/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,7 +248,11 @@ def __neg__(self) -> Self:
neg = cast("Callable[..., Column]", operator.neg)
return self._with_elementwise(neg)

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
if not strict:
msg = "`cast(..., strict=False)` is not yet implemented for spark-like backends."
raise NotImplementedError(msg)

def func(df: SparkLikeLazyFrame) -> Sequence[Column]:
spark_dtype = narwhals_to_native_dtype(
dtype, self._version, self._native_dtypes, df.native.sparkSession
Expand Down
8 changes: 6 additions & 2 deletions src/narwhals/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,11 +166,13 @@ def pipe(
"""
return function(self, *args, **kwargs)

def cast(self, dtype: IntoDType) -> Self:
def cast(self, dtype: IntoDType, *, strict: bool = True) -> Self:
"""Redefine an object's data type.

Arguments:
dtype: Data type that the object will be cast into.
strict: If `True` (default), raise an error if a value can't be cast into the
target type. If `False`, replace it with `null` instead.

Examples:
>>> import pandas as pd
Expand All @@ -188,7 +190,9 @@ def cast(self, dtype: IntoDType) -> Self:
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
"""
_validate_dtype(dtype)
return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "cast", dtype=dtype))
return self._append_node(
ExprNode(ExprKind.ELEMENTWISE, "cast", dtype=dtype, strict=strict)
)

# --- binary ---
def _with_binary(self, attr: str, other: Self | Any) -> Self:
Expand Down
Loading
Loading