fix: return null for floor division by zero on pandas-like backends

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)

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:

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,19 +143,77 @@ 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))
 
     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"]