diff --git a/src/narwhals/_pandas_like/series.py b/src/narwhals/_pandas_like/series.py index 713b250305..da402aa909 100644 --- a/src/narwhals/_pandas_like/series.py +++ b/src/narwhals/_pandas_like/series.py @@ -16,6 +16,7 @@ align_and_extract_native, binary_string_sum_fallback, broadcast_series_to_index, + floordiv_null_on_zero, get_dtype_backend, import_array_module, is_dtype_pyarrow, @@ -507,10 +508,10 @@ def __rtruediv__(self, other: Any) -> Self: return self._with_binary_right(operator.truediv, other) def __floordiv__(self, other: Any) -> Self: - return self._with_binary(operator.floordiv, other) + return self._with_binary(floordiv_null_on_zero, other) def __rfloordiv__(self, other: Any) -> Self: - return self._with_binary_right(operator.floordiv, other) + return self._with_binary_right(floordiv_null_on_zero, other) def __pow__(self, other: Any) -> Self: return self._with_binary(operator.pow, other) diff --git a/src/narwhals/_pandas_like/utils.py b/src/narwhals/_pandas_like/utils.py index 280c8a655c..34b9016210 100644 --- a/src/narwhals/_pandas_like/utils.py +++ b/src/narwhals/_pandas_like/utils.py @@ -27,6 +27,7 @@ parse_version, requires, ) +from narwhals.dependencies import is_pandas_like_series from narwhals.exceptions import ShapeError if TYPE_CHECKING: @@ -709,6 +710,39 @@ def broadcast_series_to_index( return series_class(value, index=index, dtype=native.dtype, name=native.name) +def _is_native_series(obj: object) -> TypeIs[pd.Series[Any]]: + # Narrows modin and cuDF to `pd.Series` too, matching how `_pandas_like` types them. + return is_pandas_like_series(obj) + + +def floordiv_null_on_zero( + dividend: pd.Series[Any] | float, divisor: pd.Series[Any] | float +) -> pd.Series[Any]: + """Floor-divide, with null wherever `divisor` is zero. + + Natively, integer division by zero gives `inf` (numpy-backed), `0` (nullable) + or raises (pyarrow-backed). + """ + if _is_native_series(divisor): + is_nonzero = divisor != 0 + if is_nonzero.all(): + return dividend // divisor + # A literal `1` upcasts a numpy bool divisor to object, and a nullable + # boolean one rejects it. + safe_divisor = divisor.where(is_nonzero, divisor.dtype.type(1)) + return (dividend // safe_divisor).where(is_nonzero) + assert _is_native_series(dividend) # noqa: S101 + try: + divisor_is_zero = bool(divisor == 0) + except (TypeError, ValueError): # `pd.NA == 0` and `array == 0` are ambiguous + divisor_is_zero = False + if divisor_is_zero: + all_false = type(dividend)(False, index=dividend.index) + # Every value gets masked, so this division only sets the result dtype. + return (dividend // type(divisor)(1)).where(all_false) + return dividend // divisor + + def binary_string_sum_fallback( # pragma: no cover left: pd.Series, right: Any, pdx: Any ) -> pd.Series: diff --git a/tests/expr_and_series/division_by_zero_test.py b/tests/expr_and_series/division_by_zero_test.py index 998200601a..c4ed219aad 100644 --- a/tests/expr_and_series/division_by_zero_test.py +++ b/tests/expr_and_series/division_by_zero_test.py @@ -1,15 +1,23 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import pytest import narwhals as nw -from tests.utils import POLARS_VERSION, Constructor, ConstructorEager, assert_equal_data +from tests.utils import ( + DUCKDB_VERSION, + POLARS_VERSION, + Constructor, + ConstructorEager, + assert_equal_data, +) if TYPE_CHECKING: from collections.abc import Callable + import pandas as pd + data: dict[str, list[float]] = { "int": [-2, 0, 2], "float": [-2.1, 0.0, 2.1], @@ -51,7 +59,7 @@ def test_series_floordiv_by_zero( if "polars" in str(constructor_eager) and POLARS_VERSION < (0, 20, 7): pytest.skip(reason="bug") - if df.implementation.is_pandas_like(): + if "cudf" in str(constructor_eager): request.applymarker(pytest.mark.xfail) denominator = get_denominator(df) @@ -68,7 +76,7 @@ def test_expr_floordiv_by_zero( if "polars" in str(constructor) and POLARS_VERSION < (0, 20, 7): pytest.skip(reason="bug") - if df.implementation.is_pandas_like(): + if "cudf" in str(constructor): request.applymarker(pytest.mark.xfail) result = df.select(result=nw.col("int") // denominator) @@ -120,15 +128,7 @@ def test_series_rfloordiv_by_zero( ) -> None: if "polars" in str(constructor_eager) and POLARS_VERSION < (0, 20, 7): pytest.skip(reason="bug") - if any( - x in str(constructor_eager) for x in ("pandas_pyarrow", "modin_pyarrow", "cudf") - ) or ( - any( - x in str(constructor_eager) - for x in ("pandas_nullable", "pandas_constructor", "modin_constructor") - ) - and numerator != 0 - ): + if "cudf" in str(constructor_eager): request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor_eager(data), eager_only=True) @@ -143,15 +143,7 @@ def test_expr_rfloordiv_by_zero( ) -> None: if "polars" in str(constructor) and POLARS_VERSION < (0, 20, 7): pytest.skip(reason="bug") - if any( - x in str(constructor) for x in ("pandas_pyarrow", "modin_pyarrow", "cudf") - ) or ( - any( - x in str(constructor) - for x in ("pandas_nullable", "pandas_constructor", "modin_constructor") - ) - and numerator != 0 - ): + if "cudf" in str(constructor): request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) @@ -159,3 +151,69 @@ def test_expr_rfloordiv_by_zero( result = df.select(result=numerator // nw.col("denominator")) expected = {"result": expected_floordiv} assert_equal_data(result, expected) + + +def test_floordiv_by_mixed_divisor( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if "polars" in str(constructor) and POLARS_VERSION < (0, 20, 7): + pytest.skip(reason="bug") + if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3): + pytest.skip(reason="broadcast requires `over`, which requires DuckDB 1.3.0") + if "cudf" in str(constructor): + request.applymarker(pytest.mark.xfail) + + data = {"i": [0, 1, 2, 3], "a": [7, 8, 7, None], "b": [0, 2, None, 0]} + df = nw.from_native(constructor(data)) + result = df.select( + "i", + floordiv=nw.col("a") // nw.col("b"), + rfloordiv=7 // nw.col("b"), + broadcast_dividend=nw.col("a").max() // nw.col("b"), + scalar_zero=nw.col("a") // 0, + ).sort("i") + expected = { + "i": [0, 1, 2, 3], + "floordiv": [None, 4, None, None], + "rfloordiv": [None, 3, None, None], + "broadcast_dividend": [None, 4, None, None], + "scalar_zero": [None, None, None, None], + } + assert_equal_data(result, expected) + + if df.implementation.is_pandas_like(): + # `Int64` and `int64[pyarrow]` are both `nw.Int64`, so compare native dtypes. + input_dtype = cast("pd.DataFrame", nw.to_native(df))["a"].dtype + native_result = cast("pd.DataFrame", nw.to_native(result)) + assert all( + native_result[name].dtype == input_dtype for name in expected if name != "i" + ) + + +def test_floordiv_by_null_scalar( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3): + pytest.skip(reason="broadcast requires `over`, which requires DuckDB 1.3.0") + if "pyarrow_table" in str(constructor): + # `floordiv_compat` raises on a null scalar divisor. + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor({"a": [7, 8], "b": [1, 2]})) + null_scalar = nw.when(nw.col("b") > 99).then(nw.col("b")).max() + result = df.select(nw.col("a") // null_scalar) + assert_equal_data(result, {"a": [None, None]}) + + +def test_floordiv_by_zero_keeps_dtype( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if "polars" in str(constructor) and POLARS_VERSION < (0, 20, 7): + pytest.skip(reason="bug") + if "cudf" in str(constructor): + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor({"a": [6.0, 7.0]})) + a = nw.col("a").cast(nw.Float32) + schema = df.select(by_zero=a // 0, by_two=a // 2).collect_schema() + assert schema["by_zero"] == schema["by_two"]