From a8c3ba216df5dfbfd092e51f1089eb7823771f2f Mon Sep 17 00:00:00 2001 From: Gabriel Tutui Date: Sat, 4 Jul 2020 20:29:57 +0000 Subject: [PATCH] TST: Add test for category equalness on applies (#21239) --- pandas/tests/frame/test_apply.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/pandas/tests/frame/test_apply.py b/pandas/tests/frame/test_apply.py index 114b3c0d0a3fc..3a32278e2a4b1 100644 --- a/pandas/tests/frame/test_apply.py +++ b/pandas/tests/frame/test_apply.py @@ -793,6 +793,18 @@ def test_apply_with_byte_string(self): result = df.apply(lambda x: x.astype("object")) tm.assert_frame_equal(result, expected) + @pytest.mark.parametrize("val", ["asd", 12, None, np.NaN]) + def test_apply_category_equalness(self, val): + # Check if categorical comparisons on apply, GH 21239 + df_values = ["asd", None, 12, "asd", "cde", np.NaN] + df = pd.DataFrame({"a": df_values}, dtype="category") + + result = df.a.apply(lambda x: x == val) + expected = pd.Series( + [np.NaN if pd.isnull(x) else x == val for x in df_values], name="a" + ) + tm.assert_series_equal(result, expected) + class TestInferOutputShape: # the user has supplied an opaque UDF where