Skip to content

Commit

Permalink
ENH: pd.NA comparison with time, date, timedelta
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel committed Jan 23, 2023
1 parent 1128f5e commit e237641
Show file tree
Hide file tree
Showing 6 changed files with 42 additions and 47 deletions.
9 changes: 9 additions & 0 deletions pandas/_libs/missing.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,11 @@ import numbers
from sys import maxsize

cimport cython
from cpython.datetime cimport (
date,
time,
timedelta,
)
from cython cimport Py_ssize_t

import numpy as np
Expand Down Expand Up @@ -307,6 +312,7 @@ def is_numeric_na(values: ndarray) -> ndarray:


def _create_binary_propagating_op(name, is_divmod=False):
is_cmp = name.strip("_") in ["eq", "ne", "le", "lt", "ge", "gt"]

def method(self, other):
if (other is C_NA or isinstance(other, (str, bytes))
Expand All @@ -329,6 +335,9 @@ def _create_binary_propagating_op(name, is_divmod=False):
else:
return out

elif is_cmp and isinstance(other, (date, time, timedelta)):
return NA

return NotImplemented

method.__name__ = name
Expand Down
11 changes: 9 additions & 2 deletions pandas/tests/extension/base/methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,9 @@ def test_fillna_length_mismatch(self, data_missing):
with pytest.raises(ValueError, match=msg):
data_missing.fillna(data_missing.take([1]))

# Subclasses can override if we expect e.g Sparse[bool], boolean, pyarrow[bool]
_combine_le_expected_dtype = np.dtype(bool)

def test_combine_le(self, data_repeated):
# GH 20825
# Test that combine works when doing a <= (le) comparison
Expand All @@ -268,13 +271,17 @@ def test_combine_le(self, data_repeated):
s2 = pd.Series(orig_data2)
result = s1.combine(s2, lambda x1, x2: x1 <= x2)
expected = pd.Series(
[a <= b for (a, b) in zip(list(orig_data1), list(orig_data2))]
[a <= b for (a, b) in zip(list(orig_data1), list(orig_data2))],
dtype=self._combine_le_expected_dtype,
)
self.assert_series_equal(result, expected)

val = s1.iloc[0]
result = s1.combine(val, lambda x1, x2: x1 <= x2)
expected = pd.Series([a <= val for a in list(orig_data1)])
expected = pd.Series(
[a <= val for a in list(orig_data1)],
dtype=self._combine_le_expected_dtype,
)
self.assert_series_equal(result, expected)

def test_combine_add(self, data_repeated):
Expand Down
6 changes: 1 addition & 5 deletions pandas/tests/extension/test_arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -949,11 +949,7 @@ def test_factorize(self, data_for_grouping, request):
)
super().test_factorize(data_for_grouping)

@pytest.mark.xfail(
reason="result dtype pyarrow[bool] better than expected dtype object"
)
def test_combine_le(self, data_repeated):
super().test_combine_le(data_repeated)
_combine_le_expected_dtype = "bool[pyarrow]"

def test_combine_add(self, data_repeated, request):
pa_dtype = next(data_repeated(1)).dtype.pyarrow_dtype
Expand Down
19 changes: 2 additions & 17 deletions pandas/tests/extension/test_boolean.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,8 @@ class TestReshaping(base.BaseReshapingTests):


class TestMethods(base.BaseMethodsTests):
_combine_le_expected_dtype = "boolean"

def test_factorize(self, data_for_grouping):
# override because we only have 2 unique values
labels, uniques = pd.factorize(data_for_grouping, use_na_sentinel=True)
Expand All @@ -185,23 +187,6 @@ def test_factorize(self, data_for_grouping):
tm.assert_numpy_array_equal(labels, expected_labels)
self.assert_extension_array_equal(uniques, expected_uniques)

def test_combine_le(self, data_repeated):
# override because expected needs to be boolean instead of bool dtype
orig_data1, orig_data2 = data_repeated(2)
s1 = pd.Series(orig_data1)
s2 = pd.Series(orig_data2)
result = s1.combine(s2, lambda x1, x2: x1 <= x2)
expected = pd.Series(
[a <= b for (a, b) in zip(list(orig_data1), list(orig_data2))],
dtype="boolean",
)
self.assert_series_equal(result, expected)

val = s1.iloc[0]
result = s1.combine(val, lambda x1, x2: x1 <= x2)
expected = pd.Series([a <= val for a in list(orig_data1)], dtype="boolean")
self.assert_series_equal(result, expected)

def test_searchsorted(self, data_for_sorting, as_series):
# override because we only have 2 unique values
data_for_sorting = pd.array([True, False], dtype="boolean")
Expand Down
23 changes: 1 addition & 22 deletions pandas/tests/extension/test_sparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,28 +270,7 @@ def test_fillna_frame(self, data_missing):


class TestMethods(BaseSparseTests, base.BaseMethodsTests):
def test_combine_le(self, data_repeated):
# We return a Series[SparseArray].__le__ returns a
# Series[Sparse[bool]]
# rather than Series[bool]
orig_data1, orig_data2 = data_repeated(2)
s1 = pd.Series(orig_data1)
s2 = pd.Series(orig_data2)
result = s1.combine(s2, lambda x1, x2: x1 <= x2)
expected = pd.Series(
SparseArray(
[a <= b for (a, b) in zip(list(orig_data1), list(orig_data2))],
fill_value=False,
)
)
self.assert_series_equal(result, expected)

val = s1.iloc[0]
result = s1.combine(val, lambda x1, x2: x1 <= x2)
expected = pd.Series(
SparseArray([a <= val for a in list(orig_data1)], fill_value=False)
)
self.assert_series_equal(result, expected)
_combine_le_expected_dtype = "Sparse[bool]"

def test_fillna_copy_frame(self, data_missing):
arr = data_missing.take([1, 1])
Expand Down
21 changes: 20 additions & 1 deletion pandas/tests/scalar/test_na_scalar.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
from datetime import (
date,
time,
timedelta,
)
import pickle

import numpy as np
Expand Down Expand Up @@ -67,7 +72,21 @@ def test_arithmetic_ops(all_arithmetic_functions, other):


@pytest.mark.parametrize(
"other", [NA, 1, 1.0, "a", b"a", np.int64(1), np.nan, np.bool_(True)]
"other",
[
NA,
1,
1.0,
"a",
b"a",
np.int64(1),
np.nan,
np.bool_(True),
time(0),
date(1, 2, 3),
timedelta(1),
pd.NaT,
],
)
def test_comparison_ops(comparison_op, other):
assert comparison_op(NA, other) is NA
Expand Down

0 comments on commit e237641

Please sign in to comment.