Repository: narwhals-dev/narwhals
Pull request #564 (2024-07-21 10:21:33)

feat: implement truediv for pyarrow
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.



MarcoGorelli: ```
>       assert result == [2, 1, 1 / 3]
E       assert [0.5, 1.0, 3.0] == [2, 1, 0.3333333333333333]
E         
E         At index 0 diff: 0.5 != 2
E         Use -v to get more diff

tests/series_only/arithmetic_test.py:57: AssertionError
================================================================================== short test summary info ===================================================================================
FAILED tests/series_only/arithmetic_test.py::test_truediv_same_dims[polars_series_constructor] - assert [0.5, 1.0, 3.0] == [2, 1, 0.3333333333333333]
```

hey is this a bug in Polars?

Patch:
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 67ff33de..36e5c231 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -131,6 +131,12 @@ class ArrowExpr:
     def __rfloordiv__(self, other: ArrowExpr | Any) -> Self:
         return reuse_series_implementation(self, "__rfloordiv__", other)
 
+    def __truediv__(self, other: ArrowExpr | Any) -> Self:
+        return reuse_series_implementation(self, "__truediv__", other)
+
+    def __rtruediv__(self, other: ArrowExpr | Any) -> Self:
+        return reuse_series_implementation(self, "__rtruediv__", other)
+
     def __invert__(self) -> Self:
         return reuse_series_implementation(self, "__invert__")
 
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index cb645f3d..3a3ec226 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -5,6 +5,7 @@ from typing import Any
 from typing import Iterable
 from typing import Sequence
 
+from narwhals._arrow.utils import cast_for_truediv
 from narwhals._arrow.utils import floordiv_compat
 from narwhals._arrow.utils import reverse_translate_dtype
 from narwhals._arrow.utils import translate_dtype
@@ -151,6 +152,26 @@ class ArrowSeries:
         other = validate_column_comparand(other)
         return self._from_native_series(floordiv_compat(other, ser))
 
+    def __truediv__(self, other: Any) -> Self:
+        pa = get_pyarrow()
+        pc = get_pyarrow_compute()
+        ser = self._native_series
+        other = validate_column_comparand(other)
+        if not isinstance(other, (pa.Array, pa.ChunkedArray)):
+            # scalar
+            other = pa.scalar(other)
+        return self._from_native_series(pc.divide(*cast_for_truediv(ser, other)))
+
+    def __rtruediv__(self, other: Any) -> Self:
+        pa = get_pyarrow()
+        pc = get_pyarrow_compute()
+        ser = self._native_series
+        other = validate_column_comparand(other)
+        if not isinstance(other, (pa.Array, pa.ChunkedArray)):
+            # scalar
+            other = pa.scalar(other)
+        return self._from_native_series(pc.divide(*cast_for_truediv(other, ser)))
+
     def __invert__(self) -> Self:
         pc = get_pyarrow_compute()
         return self._from_native_series(pc.invert(self._native_series))
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index 9ab594d9..e9a97ed2 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
 from typing import Any
 
 from narwhals import dtypes
@@ -7,6 +8,9 @@ from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import get_pyarrow_compute
 from narwhals.utils import isinstance_or_issubclass
 
+if TYPE_CHECKING:
+    import pyarrow as pa
+
 
 def translate_dtype(dtype: Any) -> dtypes.DType:
     pa = get_pyarrow()
@@ -223,3 +227,23 @@ def floordiv_compat(left: Any, right: Any) -> Any:
         divided = pc.divide(left, right)
         result = pc.floor(divided)
     return result
+
+
+def cast_for_truediv(
+    arrow_array: pa.ChunkedArray, pa_object: pa.Array | pa.Scalar
+) -> tuple[pa.ChunkedArray, pa.Array | pa.Scalar]:
+    # Lifted from:
+    # https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L108-L122
+    pc = get_pyarrow_compute()
+    pa = get_pyarrow()
+
+    # Ensure int / int -> float mirroring Python/Numpy behavior
+    # as pc.divide_checked(int, int) -> int
+    if pa.types.is_integer(arrow_array.type) and pa.types.is_integer(pa_object.type):
+        # GH: 56645.  # noqa: ERA001
+        # https://github.com/apache/arrow/issues/35563
+        return pc.cast(arrow_array, pa.float64(), safe=False), pc.cast(
+            pa_object, pa.float64(), safe=False
+        )
+
+    return arrow_array, pa_object


Test patch:
diff --git a/tests/expr_and_series/arithmetic_test.py b/tests/expr_and_series/arithmetic_test.py
index e43d0f30..f35a45f6 100644
--- a/tests/expr_and_series/arithmetic_test.py
+++ b/tests/expr_and_series/arithmetic_test.py
@@ -27,10 +27,7 @@ def test_arithmetic(
         request.applymarker(pytest.mark.xfail)
 
     # pyarrow case
-    if "pyarrow_table" in str(constructor) and attr in {
-        "__truediv__",
-        "__mod__",
-    }:
+    if "pyarrow_table" in str(constructor) and attr == "__mod__":
         request.applymarker(pytest.mark.xfail)
 
     data = {"a": [1, 2, 3]}
@@ -58,10 +55,7 @@ def test_right_arithmetic(
         request.applymarker(pytest.mark.xfail)
 
     # pyarrow case
-    if "table" in str(constructor) and attr in {
-        "__rtruediv__",
-        "__rmod__",
-    }:
+    if "table" in str(constructor) and attr in {"__rmod__"}:
         request.applymarker(pytest.mark.xfail)
 
     data = {"a": [1, 2, 3]}
diff --git a/tests/series_only/arithmetic_test.py b/tests/series_only/arithmetic_test.py
index 16a7e412..4668b8ad 100644
--- a/tests/series_only/arithmetic_test.py
+++ b/tests/series_only/arithmetic_test.py
@@ -40,10 +40,7 @@ def test_arithmetic(
     if "pandas_series_pyarrow" in str(constructor_series) and attr == "__mod__":
         request.applymarker(pytest.mark.xfail)
 
-    if "pyarrow_series" in str(constructor_series) and attr in {
-        "__truediv__",
-        "__mod__",
-    }:
+    if "pyarrow_series" in str(constructor_series) and attr == "__mod__":
         request.applymarker(pytest.mark.xfail)
 
     s = nw.from_native(constructor_series(data), series_only=True)
@@ -51,6 +48,18 @@ def test_arithmetic(
     assert result.to_numpy().tolist() == expected
 
 
+def test_truediv_same_dims(constructor_series: Any, request: Any) -> None:
+    if "polars" in str(constructor_series):
+        # https://github.com/pola-rs/polars/issues/17760
+        request.applymarker(pytest.mark.xfail)
+    s_left = nw.from_native(constructor_series([1, 2, 3]), series_only=True)
+    s_right = nw.from_native(constructor_series([2, 2, 1]), series_only=True)
+    result = (s_left / s_right).to_list()
+    assert result == [0.5, 1.0, 3.0]
+    result = (s_left.__rtruediv__(s_right)).to_list()
+    assert result == [2, 1, 1 / 3]
+
+
 @pytest.mark.slow()
 @given(  # type: ignore[misc]
     left=st.integers(-100, 100),