TST: Added test for query method with dtype='string' and null values pandas-dev#31516

Author: juliosanz

pandas/tests/frame/test_query_eval.py

@@ -1070,6 +1070,35 @@ def test_query_string_scalar_variable(self, parser, engine):
         r = df.query("Symbol == @symb", parser=parser, engine=engine)
         tm.assert_frame_equal(e, r)
 
+    @pytest.mark.parametrize(
+        "in_list",
+        [
+            [None, "asdf", "ghjk"],
+            ["asdf", None, "ghjk"],
+            ["asdf", "ghjk", None],
+            [None, None, "asdf"],
+            ["asdf", None, None],
+            [None, None, None],
+        ],
+    )
+    def test_query_string_null_elements(self, parser, engine, in_list):
+        # GITHUB ISSUE #31516
+        skip_if_no_pandas_parser(parser)
+        expected = {
+            i: value
+            for value, i in zip(in_list, range(len(in_list)))
+            if value is not None
+        }
+        df_expected = DataFrame({"a": expected})
+        df = DataFrame({"a": in_list})
+        res1 = df.query("a == 'asdf'", parser=parser, engine=engine)
+        res2 = df[df["a"] == "asdf"]
+        res3 = df.query("a <= 'asdf'", parser=parser, engine=engine)
+        tm.assert_frame_equal(res1, df_expected)
+        tm.assert_frame_equal(res1, res2)
+        tm.assert_frame_equal(res1, res3)
+        tm.assert_frame_equal(res2, res3)
+
 
 class TestDataFrameEvalWithFrame:
     @pytest.fixture