diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index b7f3a46f9e14..1cae1fd14f65 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -2290,15 +2290,15 @@ cdef class Array(_PandasConvertible): def __add__(self, object other): self._assert_cpu() - return _pc().call_function('add_checked', [self, other]) + return _compute_binary_op('add_checked', self, other) def __truediv__(self, object other): self._assert_cpu() - return _pc().call_function('divide_checked', [self, other]) + return _compute_binary_op('divide_checked', self, other) def __mul__(self, object other): self._assert_cpu() - return _pc().call_function('multiply_checked', [self, other]) + return _compute_binary_op('multiply_checked', self, other) def __neg__(self): self._assert_cpu() @@ -2306,31 +2306,31 @@ cdef class Array(_PandasConvertible): def __pow__(self, object other): self._assert_cpu() - return _pc().call_function('power_checked', [self, other]) + return _compute_binary_op('power_checked', self, other) def __sub__(self, object other): self._assert_cpu() - return _pc().call_function('subtract_checked', [self, other]) + return _compute_binary_op('subtract_checked', self, other) def __and__(self, object other): self._assert_cpu() - return _pc().call_function('bit_wise_and', [self, other]) + return _compute_binary_op('bit_wise_and', self, other) def __or__(self, object other): self._assert_cpu() - return _pc().call_function('bit_wise_or', [self, other]) + return _compute_binary_op('bit_wise_or', self, other) def __xor__(self, object other): self._assert_cpu() - return _pc().call_function('bit_wise_xor', [self, other]) + return _compute_binary_op('bit_wise_xor', self, other) def __lshift__(self, object other): self._assert_cpu() - return _pc().call_function('shift_left_checked', [self, other]) + return _compute_binary_op('shift_left_checked', self, other) def __rshift__(self, object other): self._assert_cpu() - return _pc().call_function('shift_right_checked', [self, other]) + return _compute_binary_op('shift_right_checked', self, other) cdef _array_like_to_pandas(obj, options, types_mapper): diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi index a6377b2bb707..0498a368b88c 100644 --- a/python/pyarrow/scalar.pxi +++ b/python/pyarrow/scalar.pxi @@ -21,6 +21,25 @@ from uuid import UUID from collections.abc import Sequence, Mapping +def _is_valid_compute_operand(value): + if isinstance(value, (Scalar, Array, ChunkedArray, RecordBatch, Table, + list, tuple)): + return True + if np is not None and isinstance(value, np.ndarray): + return True + try: + scalar(value) + except Exception: + return False + return True + + +def _compute_binary_op(func_name, left, right): + if not _is_valid_compute_operand(right): + return NotImplemented + return _pc().call_function(func_name, [left, right]) + + cdef class Scalar(_Weakrefable): """ The base class for scalars. @@ -199,37 +218,37 @@ cdef class Scalar(_Weakrefable): return _pc().call_function('abs_checked', [self]) def __add__(self, object other): - return _pc().call_function('add_checked', [self, other]) + return _compute_binary_op('add_checked', self, other) def __truediv__(self, object other): - return _pc().call_function('divide_checked', [self, other]) + return _compute_binary_op('divide_checked', self, other) def __mul__(self, object other): - return _pc().call_function('multiply_checked', [self, other]) + return _compute_binary_op('multiply_checked', self, other) def __neg__(self): return _pc().call_function('negate_checked', [self]) def __pow__(self, object other): - return _pc().call_function('power_checked', [self, other]) + return _compute_binary_op('power_checked', self, other) def __sub__(self, object other): - return _pc().call_function('subtract_checked', [self, other]) + return _compute_binary_op('subtract_checked', self, other) def __and__(self, object other): - return _pc().call_function('bit_wise_and', [self, other]) + return _compute_binary_op('bit_wise_and', self, other) def __or__(self, object other): - return _pc().call_function('bit_wise_or', [self, other]) + return _compute_binary_op('bit_wise_or', self, other) def __xor__(self, object other): - return _pc().call_function('bit_wise_xor', [self, other]) + return _compute_binary_op('bit_wise_xor', self, other) def __lshift__(self, object other): - return _pc().call_function('shift_left_checked', [self, other]) + return _compute_binary_op('shift_left_checked', self, other) def __rshift__(self, object other): - return _pc().call_function('shift_right_checked', [self, other]) + return _compute_binary_op('shift_right_checked', self, other) _NULL = NA = None diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index a103519dc5ac..38be7823b266 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -21,6 +21,7 @@ import hypothesis as h import hypothesis.strategies as st import itertools +import operator import pytest import struct import subprocess @@ -4468,3 +4469,27 @@ def test_dunders_checked_overflow(): arr ** pa.scalar(2, type=pa.int8()) with pytest.raises(pa.ArrowInvalid, match=error_match): arr / (-arr) + + +@pytest.mark.parametrize("op", [ + operator.add, + operator.sub, + operator.mul, + operator.truediv, + operator.pow, + operator.and_, + operator.or_, + operator.xor, + operator.lshift, + operator.rshift, +]) +def test_dunders_return_notimplemented_for_unknown_types(op): + # GH-49826 + class MyObj: + def __radd__(self, other): + return "reflected" + + __rsub__ = __rmul__ = __rtruediv__ = __rpow__ = __radd__ + __rand__ = __ror__ = __rxor__ = __rlshift__ = __rrshift__ = __radd__ + + assert op(pa.array([1, 2, 3]), MyObj()) == "reflected" diff --git a/python/pyarrow/tests/test_scalars.py b/python/pyarrow/tests/test_scalars.py index 08f9fcd55ce0..4ca541f4d7ce 100644 --- a/python/pyarrow/tests/test_scalars.py +++ b/python/pyarrow/tests/test_scalars.py @@ -17,6 +17,7 @@ import datetime import decimal +import operator import pytest import weakref from collections.abc import Sequence, Mapping @@ -1051,3 +1052,27 @@ def test_dunders_checked_overflow(): scl ** scl with pytest.raises(pa.ArrowInvalid, match=error_match): scl * scl + + +@pytest.mark.parametrize("op", [ + operator.add, + operator.sub, + operator.mul, + operator.truediv, + operator.pow, + operator.and_, + operator.or_, + operator.xor, + operator.lshift, + operator.rshift, +]) +def test_dunders_return_notimplemented_for_unknown_types(op): + # GH-49826 + class MyObj: + def __radd__(self, other): + return "reflected" + + __rsub__ = __rmul__ = __rtruediv__ = __rpow__ = __radd__ + __rand__ = __ror__ = __rxor__ = __rlshift__ = __rrshift__ = __radd__ + + assert op(pa.scalar(5), MyObj()) == "reflected"