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

Author: juliosanz

pandas/tests/frame/test_query_eval.py

@@ -1070,6 +1070,37 @@ 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, in_list: list):
+        # GITHUB ISSUE #31516
+        parser = "pandas"
+        engine = "python"
+        expected = {
+            i: value
+            for value, i in zip(in_list, range(len(in_list)))
+            if value == "asdf"
+        }
+        df_expected = DataFrame({"a": expected}, dtype="string")
+        df_expected.index = df_expected.index.astype("int64")
+        df = DataFrame({"a": in_list}, dtype="string")
+        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