Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions python/pyarrow/array.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -2290,47 +2290,47 @@ 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()
return _pc().call_function('negate_checked', [self])

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):
Expand Down
27 changes: 17 additions & 10 deletions python/pyarrow/scalar.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,13 @@ from uuid import UUID
from collections.abc import Sequence, Mapping


def _compute_binary_op(func_name, left, right):
try:
return _pc().call_function(func_name, [left, right])
except TypeError:
return NotImplemented


cdef class Scalar(_Weakrefable):
"""
The base class for scalars.
Expand Down Expand Up @@ -199,37 +206,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
Expand Down
25 changes: 25 additions & 0 deletions python/pyarrow/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import hypothesis as h
import hypothesis.strategies as st
import itertools
import operator
import pytest
import struct
import subprocess
Expand Down Expand Up @@ -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"
25 changes: 25 additions & 0 deletions python/pyarrow/tests/test_scalars.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

import datetime
import decimal
import operator
import pytest
import weakref
from collections.abc import Sequence, Mapping
Expand Down Expand Up @@ -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"