ENH: best attempt at NA-friendly Series comparison and boolean binary operations, GH #801

Author: wesm

pandas/core/series.py

@@ -81,6 +81,78 @@ def wrapper(self, other):
                           index=self.index, name=self.name)
     return wrapper
 
+
+def _comp_method(op, name):
+    """
+    Wrapper function for Series arithmetic operations, to avoid
+    code duplication.
+    """
+    def na_op(x, y):
+        if x.dtype == np.object_:
+            if isinstance(y, list):
+                y = lib.list_to_object_array(y)
+
+            if isinstance(y, np.ndarray):
+                result = lib.vec_compare(x, y, op)
+            else:
+                result = lib.scalar_compare(x, y, op)
+        else:
+            result = op(x, y)
+
+        return result
+
+    def wrapper(self, other):
+        from pandas.core.frame import DataFrame
+
+        if isinstance(other, Series):
+            name = _maybe_match_name(self, other)
+            return Series(na_op(self.values, other.values),
+                          index=self.index, name=name)
+        elif isinstance(other, DataFrame):
+            return NotImplemented
+        else:
+            # scalars
+            return Series(na_op(self.values, other),
+                          index=self.index, name=self.name)
+    return wrapper
+
+
+def _bool_method(op, name):
+    """
+    Wrapper function for Series arithmetic operations, to avoid
+    code duplication.
+    """
+    def na_op(x, y):
+        try:
+            result = op(x, y)
+        except TypeError:
+            if isinstance(y, list):
+                y = lib.list_to_object_array(y)
+
+            if isinstance(y, np.ndarray):
+                result = lib.vec_binop(x, y, op)
+            else:
+                result = lib.scalar_binop(x, y, op)
+
+        return result
+
+    def wrapper(self, other):
+        from pandas.core.frame import DataFrame
+
+        if isinstance(other, Series):
+            name = _maybe_match_name(self, other)
+            return Series(na_op(self.values, other.values),
+                          index=self.index, name=name)
+        elif isinstance(other, DataFrame):
+            return NotImplemented
+        else:
+            # scalars
+            return Series(na_op(self.values, other),
+                          index=self.index, name=self.name)
+    return wrapper
+
+
+
 def _radd_compat(left, right):
     radd = lambda x, y: y + x
     # GH #353, NumPy 1.5.1 workaround
@@ -97,12 +169,14 @@ def _radd_compat(left, right):
 
     return output
 
+
 def _maybe_match_name(a, b):
     name = None
     if a.name == b.name:
         name = a.name
     return name
 
+
 def _flex_method(op, name):
     doc = """
     Binary operator %s with support to substitute a fill_value for missing data
@@ -682,17 +756,17 @@ def iteritems(self, index=True):
     __rpow__ = _arith_method(lambda x, y: y ** x, '__pow__')
 
     # comparisons
-    # __gt__ = _arith_method(operator.gt, '__gt__')
-    # __ge__ = _arith_method(operator.ge, '__ge__')
-    # __lt__ = _arith_method(operator.lt, '__lt__')
-    # __le__ = _arith_method(operator.le, '__le__')
-    # __eq__ = _arith_method(operator.eq, '__eq__')
-    # __ne__ = _arith_method(operator.ne, '__ne__')
+    __gt__ = _comp_method(operator.gt, '__gt__')
+    __ge__ = _comp_method(operator.ge, '__ge__')
+    __lt__ = _comp_method(operator.lt, '__lt__')
+    __le__ = _comp_method(operator.le, '__le__')
+    __eq__ = _comp_method(operator.eq, '__eq__')
+    __ne__ = _comp_method(operator.ne, '__ne__')
 
     # binary logic
-    # __or__ = _arith_method(operator.or_, '__or__')
-    # __and__ = _arith_method(operator.and_, '__and__')
-    # __xor__ = _arith_method(operator.xor, '__xor__')
+    __or__ = _bool_method(operator.or_, '__or__')
+    __and__ = _bool_method(operator.and_, '__and__')
+    __xor__ = _bool_method(operator.xor, '__xor__')
 
     # Inplace operators
     __iadd__ = __add__
@@ -1902,7 +1976,26 @@ def isin(self, values):
         result = lib.ismember(self, value_set)
         return Series(result, self.index, name=self.name)
 
-#-------------------------------------------------------------------------------
+    def between(self, left, right, inclusive=True):
+        """
+        Return boolean Series equivalent to left <= series <= right, taking NAs
+        (if any) into account
+
+        Parameters
+        ----------
+        left : scalar
+            Left boundary
+        right : scalar
+            Right boundary
+
+        Returns
+        -------
+        is_between : Series
+            NAs, if any, will be preserved
+        """
+        pass
+
+#----------------------------------------------------------------------
 # Miscellaneous
 
     def plot(self, label=None, kind='line', use_index=True, rot=30, ax=None,

pandas/src/tseries.pyx

@@ -218,10 +218,10 @@ def array_to_datetime(ndarray[int64_t, ndim=1] arr):
 cdef double INF = <double> np.inf
 cdef double NEGINF = -INF
 
-cdef inline _checknull(object val):
+cdef inline bint _checknull(object val):
     return not np.PyArray_Check(val) and (val is None or val != val)
 
-cdef inline _checknan(object val):
+cdef inline bint _checknan(object val):
     return not np.PyArray_Check(val) and val != val
 
 cpdef checknull(object val):
@@ -232,6 +232,8 @@ cpdef checknull(object val):
     else:
         return _checknull(val)
 
+@cython.wraparound(False)
+@cython.boundscheck(False)
 def isnullobj(ndarray[object] arr):
     cdef Py_ssize_t i, n
     cdef object val
@@ -240,11 +242,11 @@ def isnullobj(ndarray[object] arr):
     n = len(arr)
     result = np.zeros(n, dtype=np.uint8)
     for i from 0 <= i < n:
-        val = arr[i]
-        if _checknull(val):
-            result[i] = 1
+        result[i] = _checknull(arr[i])
     return result.view(np.bool_)
 
+@cython.wraparound(False)
+@cython.boundscheck(False)
 def isnullobj2d(ndarray[object, ndim=2] arr):
     cdef Py_ssize_t i, j, n, m
     cdef object val
@@ -493,22 +495,117 @@ def convert_timestamps(ndarray values):
 
     return out
 
-# cdef class TypeConverter:
-#     cdef:
-#         cpython.PyTypeObject* klass_type
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def scalar_compare(ndarray[object] values, object val, object op):
+    import operator
+    cdef:
+        Py_ssize_t i, n = len(values)
+        int flag
+        object x
+
+    if op is operator.lt:
+        flag = cpython.Py_LT
+    elif op is operator.le:
+        flag = cpython.Py_LE
+    elif op is operator.gt:
+        flag = cpython.Py_GT
+    elif op is operator.ge:
+        flag = cpython.Py_GE
+    elif op is operator.eq:
+        flag = cpython.Py_EQ
+    elif op is operator.ne:
+        flag = cpython.Py_NE
+    else:
+        raise ValueError('Unrecognized operator')
+
+    result = np.empty(n, dtype=object)
+
+    for i in range(n):
+        x = values[i]
+        if _checknull(x):
+            result[i] = x
+        else:
+            result[i] = cpython.PyObject_RichCompareBool(x, val, flag)
+
+    return maybe_convert_bool(result)
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def vec_compare(ndarray[object] left, ndarray[object] right, object op):
+    import operator
+    cdef:
+        Py_ssize_t i, n = len(left)
+        int flag
+
+    if op is operator.lt:
+        flag = cpython.Py_LT
+    elif op is operator.le:
+        flag = cpython.Py_LE
+    elif op is operator.gt:
+        flag = cpython.Py_GT
+    elif op is operator.ge:
+        flag = cpython.Py_GE
+    elif op is operator.eq:
+        flag = cpython.Py_EQ
+    elif op is operator.ne:
+        flag = cpython.Py_NE
+    else:
+        raise ValueError('Unrecognized operator')
+
+    result = np.empty(n, dtype=object)
+
+    for i in range(n):
+        x = left[i]
+        y = right[i]
+        if _checknull(x):
+            result[i] = x
+        elif _checknull(y):
+            result[i] = y
+        else:
+            result[i] = cpython.PyObject_RichCompareBool(x, y, flag)
+
+    return maybe_convert_bool(result)
+
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def scalar_binop(ndarray[object] values, object val, object op):
+    cdef:
+        Py_ssize_t i, n = len(values)
+        object x
+
+    result = np.empty(n, dtype=object)
+
+    for i in range(n):
+        x = values[i]
+        if _checknull(x):
+            result[i] = x
+        else:
+            result[i] = op(x, val)
+
+    return maybe_convert_bool(result)
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def vec_binop(ndarray[object] left, ndarray[object] right, object op):
+    cdef:
+        Py_ssize_t i, n = len(left)
+
+    result = np.empty(n, dtype=object)
 
-#     cdef readonly:
-#         object factory
-#         object klass
+    for i in range(n):
+        x = left[i]
+        y = right[i]
+        if _checknull(x):
+            result[i] = x
+        elif _checknull(y):
+            result[i] = y
+        else:
+            result[i] = op(x, y)
 
-#     def __init__(self, object klass, factory):
-#         self.klass_type = (<PyObject*> klass).ob_type
-#         self.factory = factory
+    return maybe_convert_bool(result)
 
-#     def convert(self, object obj):
-#         if cpython.PyObject_TypeCheck(obj, self.klass_type):
-#             return obj
-#         return self.factory(obj)
 
 include "skiplist.pyx"
 include "groupby.pyx"

pandas/tests/test_series.py

@@ -1099,6 +1099,37 @@ def test_operators_na_handling(self):
         expected = Series(['foo_suffix', 'bar_suffix', 'baz_suffix', np.nan])
         assert_series_equal(result, expected)
 
+    def test_comparison_operators_with_nas(self):
+        from pandas import DateRange
+
+        s = Series(DateRange('1/1/2000', periods=10), dtype=object)
+        s[::2] = np.nan
+
+        # test that comparions work
+        ops = ['lt', 'le', 'gt', 'ge', 'eq', 'ne']
+        for op in ops:
+            val = s[5]
+
+            f = getattr(operator, op)
+            result = f(s, val)
+            expected = f(s.dropna(), val).reindex(s.index)
+            assert_series_equal(result, expected)
+
+            # fffffffuuuuuuuuuuuu
+            # result = f(val, s)
+            # expected = f(val, s.dropna()).reindex(s.index)
+            # assert_series_equal(result, expected)
+
+        # boolean &, |, ^ should work with object arrays and propagate NAs
+
+        ops = ['and_', 'or_', 'xor']
+        for bool_op in ops:
+            f = getattr(operator, bool_op)
+
+            result = f(s < s[9], s > s[3])
+            expected = f(s.dropna() < s[9], s.dropna() > s[3]).reindex(s.index)
+            assert_series_equal(result, expected)
+
     def test_idxmin(self):
         # test idxmin
         # _check_stat_op approach can not be used here because of isnull check.