fix overflow

Author: Rohan Jain

pandas/core/arrays/arrow/array.py

@@ -127,8 +127,10 @@ def floordiv_compat(
         left: pa.ChunkedArray | pa.Array | pa.Scalar,
         right: pa.ChunkedArray | pa.Array | pa.Scalar,
     ) -> pa.ChunkedArray:
-        divided = pc.divide(left, right)
-        if pa.types.is_integer(divided.type):
+        if pa.types.is_integer(left.type) and pa.types.is_integer(right.type):
+            # Use divide_checked to ensure cases like -9223372036854775808 // -1
+            # don't silently overflow.
+            divided = pc.divide_checked(left, right)
             # GH 56676: avoid storing intermediate calculating in floating point type.
             has_remainder = pc.not_equal(pc.multiply(divided, right), left)
             result = pc.if_else(
@@ -154,6 +156,10 @@ def floordiv_compat(
             # int8 // int64 returned int8 rather than int64.
             result = result.cast(left.type)
         else:
+            # Use divide instead of divide_checked to match numpy
+            # floordiv where divide by 0 returns infinity for floating
+            # point types.
+            divided = pc.divide(left, right)
             result = pc.floor(divided)
         return result
 

pandas/tests/extension/test_arrow.py

@@ -3261,6 +3261,23 @@ def test_arrow_floordiv_larger_divisor():
     tm.assert_series_equal(result, expected)
 
 
+def test_arrow_floordiv_integral_invalid():
+    # GH 56676
+    a = pd.Series([-9223372036854775808], dtype="int64[pyarrow]")
+    with pytest.raises(pa.lib.ArrowInvalid, match="overflow"):
+        a // -1
+    with pytest.raises(pa.lib.ArrowInvalid, match="divide by zero"):
+        a // 0
+
+
+def test_arrow_floordiv_floating_0_divisor():
+    # GH 56676
+    a = pd.Series([2], dtype="double[pyarrow]")
+    result = a // 0
+    expected = pd.Series([float("inf")], dtype="double[pyarrow]")
+    tm.assert_series_equal(result, expected)
+
+
 def test_arrow_floordiv_no_overflow():
     # GH 56676
     a = pd.Series([9223372036854775808], dtype="uint64[pyarrow]")