fix: return null for modulo by zero on pandas-like backends (#4036)

src/narwhals/_pandas_like/series.py

@@ -21,6 +21,7 @@
     import_array_module,
     is_dtype_pyarrow,
     is_pandas_or_modin,
+    mod_null_on_zero,
     narwhals_to_native_dtype,
     native_to_narwhals_dtype,
     object_native_to_narwhals_dtype,
@@ -520,10 +521,10 @@ def __rpow__(self, other: Any) -> Self:
         return self._with_binary_right(operator.pow, other)
 
     def __mod__(self, other: Any) -> Self:
-        return self._with_binary(operator.mod, other)
+        return self._with_binary(mod_null_on_zero, other)
 
     def __rmod__(self, other: Any) -> Self:
-        return self._with_binary_right(operator.mod, other)
+        return self._with_binary_right(mod_null_on_zero, other)
 
     # Unary
 

src/narwhals/_pandas_like/utils.py

@@ -715,22 +715,19 @@ def _is_native_series(obj: object) -> TypeIs[pd.Series[Any]]:
     return is_pandas_like_series(obj)
 
 
-def floordiv_null_on_zero(
-    dividend: pd.Series[Any] | float, divisor: pd.Series[Any] | float
+def _null_on_zero(
+    op: Callable[[Any, Any], Any],
+    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
+            return op(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)
+        return op(dividend, safe_divisor).where(is_nonzero)
     assert _is_native_series(dividend)  # noqa: S101
     try:
         divisor_is_zero = bool(divisor == 0)
@@ -739,8 +736,30 @@ def floordiv_null_on_zero(
     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
+        return op(dividend, type(divisor)(1)).where(all_false)
+    return op(dividend, divisor)
+
+
+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).
+    """
+    return _null_on_zero(operator.floordiv, dividend, divisor)
+
+
+def mod_null_on_zero(
+    dividend: pd.Series[Any] | float, divisor: pd.Series[Any] | float
+) -> pd.Series[Any]:
+    """Modulo, with null wherever `divisor` is zero.
+
+    Natively, modulo by zero gives `nan` (numpy-backed), `0` (nullable)
+    or raises (pyarrow-backed).
+    """
+    return _null_on_zero(operator.mod, dividend, divisor)
 
 
 def binary_string_sum_fallback(  # pragma: no cover

tests/expr_and_series/division_by_zero_test.py

@@ -25,6 +25,7 @@
 }
 expected_truediv: list[float] = [float("-inf"), float("nan"), float("inf")]
 expected_floordiv: list[float | None] = [None, None, None]
+expected_mod: list[float | None] = [None, None, None]
 
 
 @pytest.mark.parametrize("get_denominator", [lambda _: 0, lambda df: df["denominator"]])
@@ -217,3 +218,169 @@ def test_floordiv_by_zero_keeps_dtype(
     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"]
+
+
+@pytest.mark.parametrize("get_denominator", [lambda _: 0, lambda df: df["denominator"]])
+def test_series_mod_by_zero(
+    constructor_eager: ConstructorEager,
+    request: pytest.FixtureRequest,
+    get_denominator: Callable[[nw.DataFrame[Any]], int | nw.Series[Any]],
+) -> None:
+    if any(x in str(constructor_eager) for x in ("pandas_pyarrow", "modin_pyarrow")):
+        request.applymarker(
+            pytest.mark.xfail(reason="NotImplementedError: mod not implemented")
+        )
+    if "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    denominator = get_denominator(df)
+    result = {"result": df["int"] % denominator}
+    expected = {"result": expected_mod}
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize("denominator", [0, nw.lit(0), nw.col("denominator")])
+def test_expr_mod_by_zero(
+    constructor: Constructor, request: pytest.FixtureRequest, denominator: int | nw.Expr
+) -> None:
+    if any(x in str(constructor) for x in ("pandas_pyarrow", "modin_pyarrow")):
+        request.applymarker(
+            pytest.mark.xfail(reason="NotImplementedError: mod not implemented")
+        )
+    if "cudf" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor(data))
+    result = df.select(result=nw.col("int") % denominator)
+    expected = {"result": expected_mod}
+    assert_equal_data(result, expected)
+    assert_equal_data(result.select(nw.col("result").is_null().all()), {"result": [True]})
+
+
+@pytest.mark.parametrize("numerator", data["int"])
+def test_series_rmod_by_zero(
+    constructor_eager: ConstructorEager, request: pytest.FixtureRequest, numerator: float
+) -> None:
+    if any(x in str(constructor_eager) for x in ("pandas_pyarrow", "modin_pyarrow")):
+        request.applymarker(
+            pytest.mark.xfail(reason="NotImplementedError: mod not implemented")
+        )
+    if "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    result = {"result": numerator % df["denominator"]}
+    assert_equal_data(result, {"result": expected_mod})
+
+
+@pytest.mark.parametrize("numerator", data["int"])
+def test_expr_rmod_by_zero(
+    constructor: Constructor, request: pytest.FixtureRequest, numerator: float
+) -> None:
+    if any(x in str(constructor) for x in ("pandas_pyarrow", "modin_pyarrow")):
+        request.applymarker(
+            pytest.mark.xfail(reason="NotImplementedError: mod not implemented")
+        )
+    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_mod}
+    assert_equal_data(result, expected)
+
+
+def test_mod_by_mixed_divisor(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if any(x in str(constructor) for x in ("pandas_pyarrow", "modin_pyarrow")):
+        request.applymarker(
+            pytest.mark.xfail(reason="NotImplementedError: mod not implemented")
+        )
+    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",
+        mod=nw.col("a") % nw.col("b"),
+        rmod=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],
+        "mod": [None, 0, None, None],
+        "rmod": [None, 1, None, None],
+        "broadcast_dividend": [None, 0, 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_mod_by_null_scalar(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if any(x in str(constructor) for x in ("pandas_pyarrow", "modin_pyarrow")):
+        request.applymarker(
+            pytest.mark.xfail(reason="NotImplementedError: mod not implemented")
+        )
+    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):
+        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_mod_by_zero_keeps_dtype(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if any(x in str(constructor) for x in ("pandas_pyarrow", "modin_pyarrow")):
+        request.applymarker(
+            pytest.mark.xfail(reason="NotImplementedError: mod not implemented")
+        )
+    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"]
+
+
+@pytest.mark.parametrize("dtype", ["Int64", "Int32", "Int8", "UInt32"])
+def test_mod_pandas_nullable_by_zero(dtype: str) -> None:
+    pytest.importorskip("pandas")
+    import pandas as pd
+
+    values = [5, 10, 0, None]
+    s = pd.Series(values, dtype=dtype)
+    df = nw.from_native(pd.DataFrame({"a": s, "z": pd.Series([0, 0, 0, 0], dtype=dtype)}))
+
+    result_scalar = df.select(r=nw.col("a") % 0)
+    assert_equal_data(result_scalar, {"r": [None, None, None, None]})
+    assert str(nw.to_native(result_scalar)["r"].dtype) == dtype
+
+    result_col = df.select(r=nw.col("a") % nw.col("z"))
+    assert_equal_data(result_col, {"r": [None, None, None, None]})
+    assert str(nw.to_native(result_col)["r"].dtype) == dtype
+
+    result_rmod = df.select(r=7 % nw.col("z"))
+    assert_equal_data(result_rmod, {"r": [None, None, None, None]})
+    assert str(nw.to_native(result_rmod)["r"].dtype) == dtype