import operator from tokenize import TokenError import numpy as np import pytest from pandas.errors import ( NumExprClobberingError, UndefinedVariableError, ) import pandas.util._test_decorators as td import pandas as pd import pandas._testing as tm from pandas.core.computation.check import ( NUMEXPR_BLOCKED_VERSION, NUMEXPR_INSTALLED, ) skip_if_no_numexpr = pytest.mark.skipif( not NUMEXPR_INSTALLED, reason="numexpr not installed or an unsupported version" ) def skip_if_no_numexpr_fixture(): # Applied inside the engine fixture rather than as a class-level mark: the # python-engine subclasses override that fixture, but pytest collects marks # from the whole MRO, so a class mark would skip them too. if not NUMEXPR_INSTALLED: pytest.skip("numexpr not installed or an unsupported version") @pytest.fixture(params=["python", "pandas"], ids=lambda x: x) def parser(request): return request.param @pytest.fixture( params=["python", pytest.param("numexpr", marks=skip_if_no_numexpr)], ids=lambda x: x, ) def engine(request): return request.param def skip_if_no_pandas_parser(parser): if parser != "pandas": pytest.skip(f"cannot evaluate with parser={parser}") class TestCompat: @pytest.fixture def df(self): return pd.DataFrame({"A": [1, 2, 3]}) @pytest.fixture def expected1(self, df): return df[df.A > 0] @pytest.fixture def expected2(self, df): return df.A + 1 def test_query_default(self, df, expected1, expected2): # GH 12749 # this should always work, whether NUMEXPR_INSTALLED or not result = df.query("A>0") tm.assert_frame_equal(result, expected1) result = df.eval("A+1") tm.assert_series_equal(result, expected2) def test_query_None(self, df, expected1, expected2): result = df.query("A>0", engine=None) tm.assert_frame_equal(result, expected1) result = df.eval("A+1", engine=None) tm.assert_series_equal(result, expected2) def test_query_python(self, df, expected1, expected2): result = df.query("A>0", engine="python") tm.assert_frame_equal(result, expected1) result = df.eval("A+1", engine="python") tm.assert_series_equal(result, expected2) def test_query_numexpr(self, df, expected1, expected2): if NUMEXPR_INSTALLED: result = df.query("A>0", engine="numexpr") tm.assert_frame_equal(result, expected1) result = df.eval("A+1", engine="numexpr") tm.assert_series_equal(result, expected2) return if NUMEXPR_BLOCKED_VERSION is not None: msg = ( rf"numexpr {NUMEXPR_BLOCKED_VERSION} is installed, but can " r"silently return incorrect results" ) else: msg = ( r"'numexpr' is not installed or an unsupported version. " r"Cannot use engine='numexpr' for query/eval if 'numexpr' is " r"not installed" ) with pytest.raises(ImportError, match=msg): df.query("A>0", engine="numexpr") with pytest.raises(ImportError, match=msg): df.eval("A+1", engine="numexpr") class TestDataFrameEval: # smaller hits python, larger hits numexpr @pytest.mark.parametrize("n", [4, 4000]) @pytest.mark.parametrize( "op_str,op,rop", [ ("+", "__add__", "__radd__"), ("-", "__sub__", "__rsub__"), ("*", "__mul__", "__rmul__"), ("/", "__truediv__", "__rtruediv__"), ], ) def test_ops(self, op_str, op, rop, n): # tst ops and reversed ops in evaluation # GH7198 df = pd.DataFrame(1, index=range(n), columns=list("abcd")) df.iloc[0] = 2 m = df.mean() base = pd.DataFrame( # noqa: F841 np.tile(m.values, n).reshape(n, -1), columns=list("abcd") ) expected = eval(f"base {op_str} df") # ops as strings result = eval(f"m {op_str} df") tm.assert_frame_equal(result, expected) # these are commutative if op in ["+", "*"]: result = getattr(df, op)(m) tm.assert_frame_equal(result, expected) # these are not elif op in ["-", "/"]: result = getattr(df, rop)(m) tm.assert_frame_equal(result, expected) def test_dataframe_sub_numexpr_path(self): # GH7192: Note we need a large number of rows to ensure this # goes through the numexpr path df = pd.DataFrame({"A": np.random.default_rng(2).standard_normal(25000)}) df.iloc[0:5] = np.nan expected = 1 - np.isnan(df.iloc[0:25]) result = (1 - np.isnan(df)).iloc[0:25] tm.assert_frame_equal(result, expected) def test_eval_method_call_fillna(self): # GH#34045 eval must handle a chained method call such as .fillna(0) df = pd.DataFrame([1, 2], columns=["a"]) result = df.eval("((a - 1) / (a - 1)).fillna(0)") expected = pd.Series([0.0, 1.0], name="a") tm.assert_series_equal(result, expected) def test_query_method_call_isnull(self, engine): # GH#34251 query must support a method call such as name.isnull(). # The crash was numexpr-engine-specific, so exercise both engines # explicitly (the default silently falls back to python without numexpr) df = pd.DataFrame({"id": [0, 1, 2], "name": ["a", None, None]}) result = df.query("name.isnull()", engine=engine) expected = df.iloc[[1, 2]] tm.assert_frame_equal(result, expected) def test_query_non_str(self): # GH 11485 df = pd.DataFrame({"A": [1, 2, 3], "B": ["a", "b", "b"]}) msg = "expr must be a string to be evaluated" with pytest.raises(ValueError, match=msg): df.query(lambda x: x.B == "b") with pytest.raises(ValueError, match=msg): df.query(111) def test_query_empty_string(self): # GH 13139 df = pd.DataFrame({"A": [1, 2, 3]}) msg = "expr cannot be an empty string" with pytest.raises(ValueError, match=msg): df.query("") def test_query_duplicate_column_name(self, engine, parser): df = pd.DataFrame({"A": range(3), "B": range(3), "C": range(3)}).rename( columns={"B": "A"} ) res = df.query("C == 1", engine=engine, parser=parser) expect = pd.DataFrame([[1, 1, 1]], columns=["A", "A", "C"], index=[1]) tm.assert_frame_equal(res, expect) def test_query_duplicate_column_name_referenced(self, engine, parser): # GH#65588 the eval fix made a duplicated label resolve to a DataFrame; # query cannot use that as a row mask, so it raises instead of masking. df = pd.DataFrame({"A": range(3), "B": range(10, 13), "C": range(3)}).rename( columns={"B": "A"} ) msg = "expr referenced a duplicated column label" with pytest.raises(ValueError, match=msg): df.query("A == 1", engine=engine, parser=parser) # every column duplicated -> the mask needs no reindexing, so nothing # but this check stops query from silently returning a NaN-masked frame df2 = pd.DataFrame({"A": range(3), "B": range(10, 13)}).rename( columns={"B": "A"} ) with pytest.raises(ValueError, match=msg): df2.query("A == 1", engine=engine, parser=parser) def test_query_duplicate_column_name_local_not_shadowed(self, engine, parser): # GH#65588 the duplicate-label recorder must not capture an @local whose # name matches a duplicated column: Scope.swapkey rewrites the local by # writing into the first resolver that *contains* the name skip_if_no_pandas_parser(parser) df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["AA", "AA", "B"]) AA = 4 # noqa: F841 result = df.query("B > @AA", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[[1]]) @pytest.mark.parametrize("name", ["Timestamp", "datetime", "list", "tuple"]) def test_query_index_name_shadowing_default_global(self, name, engine, parser): # GH#34958, GH#35595 an index level wins over the type-valued # DEFAULT_GLOBALS; test_eval.py::test_eval_no_support_column_name # covers the column side df = pd.DataFrame({"b": [3, 4]}, index=pd.Index([1, 20], name=name)) result = df.query(f"{name} < 10", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[[0]]) def test_query_eval_local_class_attribute(self, engine, parser): # GH#48694 an @-prefixed name is local even when it refers to a class skip_if_no_pandas_parser(parser) class A: a = 2 df = pd.DataFrame({"x": [1, 2, 3]}) result = df.query("x == @A.a", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[[1]]) result = df.eval("x == @A.a", engine=engine, parser=parser) tm.assert_series_equal(result, df["x"] == 2) def test_query_duplicate_column_name_cleaned_name_collision(self, engine, parser): # GH#65588 clean_column_name is not injective, so the recorder has to # ask which label the cleaned name actually resolves to (the last one # wins) rather than whether any label sharing it is duplicated skip_if_no_pandas_parser(parser) # `1` resolves to the unique str column, so this is not ambiguous df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=[1, 1, "1"]) result = df.query("`1` == 3", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[[0]]) # reversed, `1` resolves to the duplicated int label instead df2 = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=[1, "1", 1]) with pytest.raises(ValueError, match="referenced a duplicated column label"): df2.query("`1` == 3", engine=engine, parser=parser) def test_query_duplicate_column_name_reduced_to_mask(self, engine, parser): # GH#65588 a duplicated label is only a problem when the expression # leaves the result 2-D; reducing it back to a row mask is valid and # worked before this guard existed df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) df.columns = ["a", "a", "c"] result = df.query("(a == 1).any(axis=1)", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[[0]]) result = df.query("a.max(axis=1) > 4", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[[1, 2]]) def test_query_duplicate_column_name_whole_frame_mask(self, engine, parser): # GH#65588 a 2-D mask that does not come from a duplicated label is # still a valid query, whether or not it references a unique label skip_if_no_pandas_parser(parser) df = pd.DataFrame(np.arange(-3, 3).reshape(2, 3), columns=["a", "a", "b"]) result = df.query("@df > 0", engine=engine, parser=parser) tm.assert_frame_equal(result, df[df > 0]) # referencing the unique label exercises the recorder's negative branch expected = df[df > df["b"].min()] assert expected.notna().any().any() # not vacuously all-NaN result = df.query("@df > b.min()", engine=engine, parser=parser) tm.assert_frame_equal(result, expected) def test_query_duplicate_index_label(self, engine, parser): # GH#51815 aligning the terms against a frame whose index has duplicate # labels raised "cannot reindex on an axis with duplicate labels" df = pd.Series([1, 2, 3], index=[1, 1, 2], name="col").to_frame() result = df.query("(index == 1) & (col == 2)", engine=engine, parser=parser) expected = pd.DataFrame({"col": [2]}, index=[1]) tm.assert_frame_equal(result, expected) def test_query_datetime_compared_to_string_no_warning(self, engine, parser): # GH#57028 comparing a datetime64 column to a string warned about the # behavior of 'isin', which the expression does not use. # Comparing against a string literal is rewritten to a membership op, # which the python parser does not implement. skip_if_no_pandas_parser(parser) df = pd.DataFrame({"sent": [pd.Timestamp("2024-01-14"), pd.NaT, pd.NaT]}) with tm.assert_produces_warning(None): result = df.query("sent == ''", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[:0]) def test_query_datetime_in_strings_no_warning(self, engine, parser): # GH#57028 the `in` operator does go through isin, and strings no longer # match datetime64 values there, but neither should warn skip_if_no_pandas_parser(parser) df = pd.DataFrame({"sent": [pd.Timestamp("2024-01-14"), pd.NaT, pd.NaT]}) with tm.assert_produces_warning(None): result = df.query("sent in ['2024-01-14']", engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[:0]) def test_eval_duplicate_column_name(self, engine, parser): # GH#65588 df = pd.DataFrame({"a": range(3), "b": range(10, 13), "c": range(3)}).rename( columns={"b": "a"} ) result_frame = df.eval("a == 1", engine=engine, parser=parser) result_func = pd.eval("a == 1", resolvers=(df,), engine=engine, parser=parser) expected = pd.DataFrame( [[False, False], [True, False], [False, False]], columns=["a", "a"], ) tm.assert_frame_equal(result_frame, expected) tm.assert_frame_equal(result_func, expected) def test_eval_resolvers_as_list(self): # GH 14095 df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 2)), columns=list("ab") ) dict1 = {"a": 1} dict2 = {"b": 2} assert df.eval("a + b", resolvers=[dict1, dict2]) == dict1["a"] + dict2["b"] assert pd.eval("a + b", resolvers=[dict1, dict2]) == dict1["a"] + dict2["b"] def test_eval_resolvers_combined(self): # GH 34966 df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 2)), columns=list("ab") ) dict1 = {"c": 2} # Both input and default index/column resolvers should be usable result = df.eval("a + b * c", resolvers=[dict1]) expected = df["a"] + df["b"] * dict1["c"] tm.assert_series_equal(result, expected) def test_eval_object_dtype_binop(self): # GH#24883 df = pd.DataFrame({"a1": ["Y", "N"]}) res = df.eval("c = ((a1 == 'Y') & True)") expected = pd.DataFrame({"a1": ["Y", "N"], "c": [True, False]}) tm.assert_frame_equal(res, expected) def test_using_numpy(self, engine, parser): # GH 58041 skip_if_no_pandas_parser(parser) df = pd.Series([0.2, 1.5, 2.8], name="a").to_frame() res = df.eval("@np.floor(a)", engine=engine, parser=parser) expected = np.floor(df["a"]) tm.assert_series_equal(expected, res) def test_using_local_function_with_scalar_arg(self, engine, parser): # GH#22649 the bare-name branch of visit_Call; test_using_numpy covers # the same fix through the attribute branch skip_if_no_pandas_parser(parser) around = np.around # noqa: F841 df = pd.Series([1.12, 2.76], name="a").to_frame() result = df.eval("@around(a, 1)", engine=engine, parser=parser) expected = np.around(df["a"], 1) tm.assert_series_equal(result, expected) def test_eval_simple(self, engine, parser): df = pd.Series([0.2, 1.5, 2.8], name="a").to_frame() res = df.eval("a", engine=engine, parser=parser) expected = df["a"] tm.assert_series_equal(expected, res) def test_extension_array_eval(self, engine, parser, request): # GH#58748 if engine == "numexpr": mark = pytest.mark.xfail( reason="numexpr does not support extension array dtypes" ) request.applymarker(mark) df = pd.DataFrame({"a": pd.array([1, 2, 3]), "b": pd.array([4, 5, 6])}) result = df.eval("a / b", engine=engine, parser=parser) expected = pd.Series(pd.array([0.25, 0.40, 0.50])) tm.assert_series_equal(result, expected) def test_complex_eval(self, engine, parser): # GH#21374 df = pd.DataFrame({"a": [1 + 2j], "b": [1 + 1j]}) result = df.eval("a/b", engine=engine, parser=parser) expected = pd.Series([1.5 + 0.5j]) tm.assert_series_equal(result, expected) def test_query_period_dtype(self, engine, parser): # GH#35247 df = pd.DataFrame( { "year": pd.period_range("2018", periods=5, freq="Y"), "val": range(5), } ) result = df.query("year > '2020'", engine=engine, parser=parser) expected = df[df["year"] > "2020"] tm.assert_frame_equal(result, expected) def test_query_interval_dtype(self, engine, parser): # GH#35247 idx = pd.IntervalIndex.from_breaks(range(6)) df = pd.DataFrame({"a": idx, "b": idx}) result = df.query("a == b", engine=engine, parser=parser) tm.assert_frame_equal(result, df) class TestDataFrameQueryWithMultiIndex: def test_query_with_named_multiindex(self, parser, engine): skip_if_no_pandas_parser(parser) a = np.random.default_rng(2).choice(["red", "green"], size=10) b = np.random.default_rng(2).choice(["eggs", "ham"], size=10) index = pd.MultiIndex.from_arrays([a, b], names=["color", "food"]) df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 2)), index=index ) ind = pd.Series( df.index.get_level_values("color").values, index=index, name="color" ) # equality res1 = df.query('color == "red"', parser=parser, engine=engine) res2 = df.query('"red" == color', parser=parser, engine=engine) exp = df[ind == "red"] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # inequality res1 = df.query('color != "red"', parser=parser, engine=engine) res2 = df.query('"red" != color', parser=parser, engine=engine) exp = df[ind != "red"] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # list equality (really just set membership) res1 = df.query('color == ["red"]', parser=parser, engine=engine) res2 = df.query('["red"] == color', parser=parser, engine=engine) exp = df[ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) res1 = df.query('color != ["red"]', parser=parser, engine=engine) res2 = df.query('["red"] != color', parser=parser, engine=engine) exp = df[~ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # in/not in ops res1 = df.query('["red"] in color', parser=parser, engine=engine) res2 = df.query('"red" in color', parser=parser, engine=engine) exp = df[ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) res1 = df.query('["red"] not in color', parser=parser, engine=engine) res2 = df.query('"red" not in color', parser=parser, engine=engine) exp = df[~ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) def test_query_with_unnamed_multiindex(self, parser, engine): skip_if_no_pandas_parser(parser) a = np.random.default_rng(2).choice(["red", "green"], size=10) b = np.random.default_rng(2).choice(["eggs", "ham"], size=10) index = pd.MultiIndex.from_arrays([a, b]) df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 2)), index=index ) ind = pd.Series(df.index.get_level_values(0).values, index=index) res1 = df.query('ilevel_0 == "red"', parser=parser, engine=engine) res2 = df.query('"red" == ilevel_0', parser=parser, engine=engine) exp = df[ind == "red"] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # inequality res1 = df.query('ilevel_0 != "red"', parser=parser, engine=engine) res2 = df.query('"red" != ilevel_0', parser=parser, engine=engine) exp = df[ind != "red"] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # list equality (really just set membership) res1 = df.query('ilevel_0 == ["red"]', parser=parser, engine=engine) res2 = df.query('["red"] == ilevel_0', parser=parser, engine=engine) exp = df[ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) res1 = df.query('ilevel_0 != ["red"]', parser=parser, engine=engine) res2 = df.query('["red"] != ilevel_0', parser=parser, engine=engine) exp = df[~ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # in/not in ops res1 = df.query('["red"] in ilevel_0', parser=parser, engine=engine) res2 = df.query('"red" in ilevel_0', parser=parser, engine=engine) exp = df[ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) res1 = df.query('["red"] not in ilevel_0', parser=parser, engine=engine) res2 = df.query('"red" not in ilevel_0', parser=parser, engine=engine) exp = df[~ind.isin(["red"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # ## LEVEL 1 ind = pd.Series(df.index.get_level_values(1).values, index=index) res1 = df.query('ilevel_1 == "eggs"', parser=parser, engine=engine) res2 = df.query('"eggs" == ilevel_1', parser=parser, engine=engine) exp = df[ind == "eggs"] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # inequality res1 = df.query('ilevel_1 != "eggs"', parser=parser, engine=engine) res2 = df.query('"eggs" != ilevel_1', parser=parser, engine=engine) exp = df[ind != "eggs"] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # list equality (really just set membership) res1 = df.query('ilevel_1 == ["eggs"]', parser=parser, engine=engine) res2 = df.query('["eggs"] == ilevel_1', parser=parser, engine=engine) exp = df[ind.isin(["eggs"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) res1 = df.query('ilevel_1 != ["eggs"]', parser=parser, engine=engine) res2 = df.query('["eggs"] != ilevel_1', parser=parser, engine=engine) exp = df[~ind.isin(["eggs"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) # in/not in ops res1 = df.query('["eggs"] in ilevel_1', parser=parser, engine=engine) res2 = df.query('"eggs" in ilevel_1', parser=parser, engine=engine) exp = df[ind.isin(["eggs"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) res1 = df.query('["eggs"] not in ilevel_1', parser=parser, engine=engine) res2 = df.query('"eggs" not in ilevel_1', parser=parser, engine=engine) exp = df[~ind.isin(["eggs"])] tm.assert_frame_equal(res1, exp) tm.assert_frame_equal(res2, exp) def test_query_with_partially_named_multiindex(self, parser, engine): skip_if_no_pandas_parser(parser) a = np.random.default_rng(2).choice(["red", "green"], size=10) b = np.arange(10) index = pd.MultiIndex.from_arrays([a, b]) index.names = [None, "rating"] df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 2)), index=index ) res = df.query("rating == 1", parser=parser, engine=engine) ind = pd.Series( df.index.get_level_values("rating").values, index=index, name="rating" ) exp = df[ind == 1] tm.assert_frame_equal(res, exp) res = df.query("rating != 1", parser=parser, engine=engine) ind = pd.Series( df.index.get_level_values("rating").values, index=index, name="rating" ) exp = df[ind != 1] tm.assert_frame_equal(res, exp) res = df.query('ilevel_0 == "red"', parser=parser, engine=engine) ind = pd.Series(df.index.get_level_values(0).values, index=index) exp = df[ind == "red"] tm.assert_frame_equal(res, exp) res = df.query('ilevel_0 != "red"', parser=parser, engine=engine) ind = pd.Series(df.index.get_level_values(0).values, index=index) exp = df[ind != "red"] tm.assert_frame_equal(res, exp) def test_query_multiindex_get_index_resolvers(self): df = pd.DataFrame( np.ones((10, 3)), index=pd.MultiIndex.from_arrays( [range(10) for _ in range(2)], names=["spam", "eggs"] ), ) resolvers = df._get_index_resolvers() def to_series(mi, level): level_values = mi.get_level_values(level) s = level_values.to_series() s.index = mi return s col_series = df.columns.to_series() expected = { "index": df.index, "columns": col_series, "spam": to_series(df.index, "spam"), "eggs": to_series(df.index, "eggs"), "clevel_0": col_series, } for k, v in resolvers.items(): if isinstance(v, pd.Index): assert v.is_(expected[k]) elif isinstance(v, pd.Series): tm.assert_series_equal(v, expected[k]) else: raise AssertionError("object must be a Series or Index") class TestDataFrameQueryNumExprPandas: @pytest.fixture def engine(self): skip_if_no_numexpr_fixture() return "numexpr" @pytest.fixture def parser(self): return "pandas" def test_date_query_with_attribute_access(self, engine, parser): skip_if_no_pandas_parser(parser) df = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=5) df["dates2"] = pd.date_range("1/1/2013", periods=5) df["dates3"] = pd.date_range("1/1/2014", periods=5) res = df.query( "@df.dates1 < 20130101 < @df.dates3", engine=engine, parser=parser ) expec = df[(df.dates1 < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_query_no_attribute_access(self, engine, parser): df = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=5) df["dates2"] = pd.date_range("1/1/2013", periods=5) df["dates3"] = pd.date_range("1/1/2014", periods=5) res = df.query("dates1 < 20130101 < dates3", engine=engine, parser=parser) expec = df[(df.dates1 < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_query_with_NaT(self, engine, parser): n = 10 df = pd.DataFrame(np.random.default_rng(2).standard_normal((n, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=n) df["dates2"] = pd.date_range("1/1/2013", periods=n) df["dates3"] = pd.date_range("1/1/2014", periods=n) df.loc[np.random.default_rng(2).random(n) > 0.5, "dates1"] = pd.NaT df.loc[np.random.default_rng(2).random(n) > 0.5, "dates3"] = pd.NaT res = df.query("dates1 < 20130101 < dates3", engine=engine, parser=parser) expec = df[(df.dates1 < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_index_query(self, engine, parser): n = 10 df = pd.DataFrame(np.random.default_rng(2).standard_normal((n, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=n) df["dates3"] = pd.date_range("1/1/2014", periods=n) df = df.set_index("dates1") res = df.query("index < 20130101 < dates3", engine=engine, parser=parser) expec = df[(df.index < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_index_query_with_NaT(self, engine, parser): n = 10 # Cast to object to avoid implicit cast when setting entry to pd.NaT below df = pd.DataFrame(np.random.default_rng(2).standard_normal((n, 3))).astype( {0: object} ) df["dates1"] = pd.date_range("1/1/2012", periods=n) df["dates3"] = pd.date_range("1/1/2014", periods=n) df.iloc[0, 0] = pd.NaT df = df.set_index("dates1") res = df.query("index < 20130101 < dates3", engine=engine, parser=parser) expec = df[(df.index < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_index_query_with_NaT_duplicates(self, engine, parser): n = 10 d = {} d["dates1"] = pd.date_range("1/1/2012", periods=n) d["dates3"] = pd.date_range("1/1/2014", periods=n) df = pd.DataFrame(d) df.loc[np.random.default_rng(2).random(n) > 0.5, "dates1"] = pd.NaT df = df.set_index("dates1") res = df.query("dates1 < 20130101 < dates3", engine=engine, parser=parser) expec = df[(df.index.to_series() < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_query_with_non_date(self, engine, parser): n = 10 df = pd.DataFrame( { "dates": pd.date_range("1/1/2012", periods=n, unit="ns"), "nondate": np.arange(n), } ) result = df.query("dates == nondate", parser=parser, engine=engine) assert len(result) == 0 result = df.query("dates != nondate", parser=parser, engine=engine) tm.assert_frame_equal(result, df) msg = r"Invalid comparison between dtype=datetime64\[ns\] and ndarray" for op in ["<", ">", "<=", ">="]: with pytest.raises(TypeError, match=msg): df.query(f"dates {op} nondate", parser=parser, engine=engine) def test_query_syntax_error(self, engine, parser): df = pd.DataFrame({"i": range(10), "+": range(3, 13), "r": range(4, 14)}) msg = "invalid syntax" with pytest.raises(SyntaxError, match=msg): df.query("i - +", engine=engine, parser=parser) def test_query_scope(self, engine, parser): skip_if_no_pandas_parser(parser) df = pd.DataFrame( np.random.default_rng(2).standard_normal((20, 2)), columns=list("ab") ) a, b = 1, 2 # noqa: F841 res = df.query("a > b", engine=engine, parser=parser) expected = df[df.a > df.b] tm.assert_frame_equal(res, expected) res = df.query("@a > b", engine=engine, parser=parser) expected = df[a > df.b] tm.assert_frame_equal(res, expected) # no local variable c with pytest.raises( UndefinedVariableError, match="local variable 'c' is not defined" ): df.query("@a > b > @c", engine=engine, parser=parser) # no column named 'c' with pytest.raises(UndefinedVariableError, match="name 'c' is not defined"): df.query("@a > b > c", engine=engine, parser=parser) def test_query_doesnt_pickup_local(self, engine, parser): n = m = 10 df = pd.DataFrame( np.random.default_rng(2).integers(m, size=(n, 3)), columns=list("abc") ) # we don't pick up the local 'sin' with pytest.raises(UndefinedVariableError, match="name 'sin' is not defined"): df.query("sin > 5", engine=engine, parser=parser) def test_query_builtin(self, engine, parser): n = m = 10 df = pd.DataFrame( np.random.default_rng(2).integers(m, size=(n, 3)), columns=list("abc") ) df.index.name = "sin" msg = "Variables in expression.+" with pytest.raises(NumExprClobberingError, match=msg): df.query("sin > 5", engine=engine, parser=parser) def test_query(self, engine, parser): df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 3)), columns=["a", "b", "c"] ) tm.assert_frame_equal( df.query("a < b", engine=engine, parser=parser), df[df.a < df.b] ) tm.assert_frame_equal( df.query("a + b > b * c", engine=engine, parser=parser), df[df.a + df.b > df.b * df.c], ) def test_query_index_with_name(self, engine, parser): df = pd.DataFrame( np.random.default_rng(2).integers(10, size=(10, 3)), index=pd.Index(range(10), name="blob"), columns=["a", "b", "c"], ) res = df.query("(blob < 5) & (a < b)", engine=engine, parser=parser) expec = df[(df.index < 5) & (df.a < df.b)] tm.assert_frame_equal(res, expec) res = df.query("blob < b", engine=engine, parser=parser) expec = df[df.index < df.b] tm.assert_frame_equal(res, expec) def test_query_index_without_name(self, engine, parser): df = pd.DataFrame( np.random.default_rng(2).integers(10, size=(10, 3)), index=range(10), columns=["a", "b", "c"], ) # "index" should refer to the index res = df.query("index < b", engine=engine, parser=parser) expec = df[df.index < df.b] tm.assert_frame_equal(res, expec) # test against a scalar res = df.query("index < 5", engine=engine, parser=parser) expec = df[df.index < 5] tm.assert_frame_equal(res, expec) def test_nested_scope(self, engine, parser): skip_if_no_pandas_parser(parser) df = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) df2 = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) expected = df[(df > 0) & (df2 > 0)] result = df.query("(@df > 0) & (@df2 > 0)", engine=engine, parser=parser) tm.assert_frame_equal(result, expected) result = pd.eval("df[df > 0 and df2 > 0]", engine=engine, parser=parser) tm.assert_frame_equal(result, expected) result = pd.eval( "df[df > 0 and df2 > 0 and df[df > 0] > 0]", engine=engine, parser=parser ) expected = df[(df > 0) & (df2 > 0) & (df[df > 0] > 0)] tm.assert_frame_equal(result, expected) result = pd.eval("df[(df>0) & (df2>0)]", engine=engine, parser=parser) expected = df.query("(@df>0) & (@df2>0)", engine=engine, parser=parser) tm.assert_frame_equal(result, expected) def test_nested_raises_on_local_self_reference(self, engine, parser): df = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) # can't reference ourself b/c we're a local so @ is necessary with pytest.raises(UndefinedVariableError, match="name 'df' is not defined"): df.query("df > 0", engine=engine, parser=parser) def test_local_syntax(self, engine, parser): skip_if_no_pandas_parser(parser) df = pd.DataFrame( np.random.default_rng(2).standard_normal((100, 10)), columns=list("abcdefghij"), ) b = 1 expect = df[df.a < b] result = df.query("a < @b", engine=engine, parser=parser) tm.assert_frame_equal(result, expect) expect = df[df.a < df.b] result = df.query("a < b", engine=engine, parser=parser) tm.assert_frame_equal(result, expect) def test_chained_cmp_and_in(self, engine, parser): skip_if_no_pandas_parser(parser) cols = list("abc") df = pd.DataFrame( np.random.default_rng(2).standard_normal((100, len(cols))), columns=cols ) res = df.query( "a < b < c and a not in b not in c", engine=engine, parser=parser ) ind = (df.a < df.b) & (df.b < df.c) & ~df.b.isin(df.a) & ~df.c.isin(df.b) expec = df[ind] tm.assert_frame_equal(res, expec) def test_local_variable_with_in(self, engine, parser): skip_if_no_pandas_parser(parser) a = pd.Series(np.random.default_rng(2).integers(3, size=15), name="a") b = pd.Series(np.random.default_rng(2).integers(10, size=15), name="b") df = pd.DataFrame({"a": a, "b": b}) expected = df.loc[(df.b - 1).isin(a)] result = df.query("b - 1 in a", engine=engine, parser=parser) tm.assert_frame_equal(expected, result) b = pd.Series(np.random.default_rng(2).integers(10, size=15), name="b") expected = df.loc[(b - 1).isin(a)] result = df.query("@b - 1 in a", engine=engine, parser=parser) tm.assert_frame_equal(expected, result) def test_at_inside_string(self, engine, parser): skip_if_no_pandas_parser(parser) c = 1 # noqa: F841 df = pd.DataFrame({"a": ["a", "a", "b", "b", "@c", "@c"]}) result = df.query('a == "@c"', engine=engine, parser=parser) expected = df[df.a == "@c"] tm.assert_frame_equal(result, expected) def test_query_undefined_local(self, engine, parser): skip_if_no_pandas_parser(parser) df = pd.DataFrame(np.random.default_rng(2).random((10, 2)), columns=list("ab")) with pytest.raises( UndefinedVariableError, match="local variable 'c' is not defined" ): df.query("a == @c", engine=engine, parser=parser) def test_index_resolvers_come_after_columns_with_the_same_name( self, engine, parser ): n = 1 # noqa: F841 a = np.r_[20:101:20] df = pd.DataFrame( {"index": a, "b": np.random.default_rng(2).standard_normal(a.size)} ) df.index.name = "index" result = df.query("index > 5", engine=engine, parser=parser) expected = df[df["index"] > 5] tm.assert_frame_equal(result, expected) df = pd.DataFrame( {"index": a, "b": np.random.default_rng(2).standard_normal(a.size)} ) result = df.query("ilevel_0 > 5", engine=engine, parser=parser) expected = df.loc[df.index[df.index > 5]] tm.assert_frame_equal(result, expected) df = pd.DataFrame( {"a": a, "b": np.random.default_rng(2).standard_normal(a.size)} ) df.index.name = "a" result = df.query("a > 5", engine=engine, parser=parser) expected = df[df.a > 5] tm.assert_frame_equal(result, expected) result = df.query("index > 5", engine=engine, parser=parser) expected = df.loc[df.index[df.index > 5]] tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("op, f", [["==", operator.eq], ["!=", operator.ne]]) def test_inf(self, op, f, engine, parser): n = 10 df = pd.DataFrame( { "a": np.random.default_rng(2).random(n), "b": np.random.default_rng(2).random(n), } ) df.loc[::2, 0] = np.inf q = f"a {op} inf" expected = df[f(df.a, np.inf)] result = df.query(q, engine=engine, parser=parser) tm.assert_frame_equal(result, expected) def test_check_tz_aware_index_query(self, tz_aware_fixture, engine, parser): # https://github.com/pandas-dev/pandas/issues/29463 tz = tz_aware_fixture df_index = pd.date_range( start="2019-01-01", freq="1D", periods=10, tz=tz, name="time" ) expected = pd.DataFrame(index=df_index) df = pd.DataFrame(index=df_index) result = df.query( '"2018-01-03 00:00:00+00" < time', engine=engine, parser=parser ) tm.assert_frame_equal(result, expected) expected = pd.DataFrame(df_index) result = df.reset_index().query( '"2018-01-03 00:00:00+00" < time', engine=engine, parser=parser ) tm.assert_frame_equal(result, expected) def test_method_calls_in_query(self, engine, parser): # https://github.com/pandas-dev/pandas/issues/22435 n = 10 df = pd.DataFrame( { "a": 2 * np.random.default_rng(2).random(n), "b": np.random.default_rng(2).random(n), } ) expected = df[df["a"].astype("int") == 0] result = df.query("a.astype('int') == 0", engine=engine, parser=parser) tm.assert_frame_equal(result, expected) df = pd.DataFrame( { "a": np.where( np.random.default_rng(2).random(n) < 0.5, np.nan, np.random.default_rng(2).standard_normal(n), ), "b": np.random.default_rng(2).standard_normal(n), } ) expected = df[df["a"].notnull()] result = df.query("a.notnull()", engine=engine, parser=parser) tm.assert_frame_equal(result, expected) class TestDataFrameQueryNumExprPython(TestDataFrameQueryNumExprPandas): @pytest.fixture def engine(self): skip_if_no_numexpr_fixture() return "numexpr" @pytest.fixture def parser(self): return "python" def test_date_query_no_attribute_access(self, engine, parser): df = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=5) df["dates2"] = pd.date_range("1/1/2013", periods=5) df["dates3"] = pd.date_range("1/1/2014", periods=5) res = df.query( "(dates1 < 20130101) & (20130101 < dates3)", engine=engine, parser=parser ) expec = df[(df.dates1 < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_query_with_NaT(self, engine, parser): n = 10 df = pd.DataFrame(np.random.default_rng(2).standard_normal((n, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=n) df["dates2"] = pd.date_range("1/1/2013", periods=n) df["dates3"] = pd.date_range("1/1/2014", periods=n) df.loc[np.random.default_rng(2).random(n) > 0.5, "dates1"] = pd.NaT df.loc[np.random.default_rng(2).random(n) > 0.5, "dates3"] = pd.NaT res = df.query( "(dates1 < 20130101) & (20130101 < dates3)", engine=engine, parser=parser ) expec = df[(df.dates1 < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_index_query(self, engine, parser): n = 10 df = pd.DataFrame(np.random.default_rng(2).standard_normal((n, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=n) df["dates3"] = pd.date_range("1/1/2014", periods=n) df = df.set_index("dates1") res = df.query( "(index < 20130101) & (20130101 < dates3)", engine=engine, parser=parser ) expec = df[(df.index < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_index_query_with_NaT(self, engine, parser): n = 10 # Cast to object to avoid implicit cast when setting entry to pd.NaT below df = pd.DataFrame(np.random.default_rng(2).standard_normal((n, 3))).astype( {0: object} ) df["dates1"] = pd.date_range("1/1/2012", periods=n) df["dates3"] = pd.date_range("1/1/2014", periods=n) df.iloc[0, 0] = pd.NaT df = df.set_index("dates1") res = df.query( "(index < 20130101) & (20130101 < dates3)", engine=engine, parser=parser ) expec = df[(df.index < "20130101") & ("20130101" < df.dates3)] tm.assert_frame_equal(res, expec) def test_date_index_query_with_NaT_duplicates(self, engine, parser): n = 10 df = pd.DataFrame(np.random.default_rng(2).standard_normal((n, 3))) df["dates1"] = pd.date_range("1/1/2012", periods=n) df["dates3"] = pd.date_range("1/1/2014", periods=n) df.loc[np.random.default_rng(2).random(n) > 0.5, "dates1"] = pd.NaT df = df.set_index("dates1") msg = r"'BoolOp' nodes are not implemented" with pytest.raises(NotImplementedError, match=msg): df.query("index < 20130101 < dates3", engine=engine, parser=parser) def test_nested_scope(self, engine, parser): # smoke test x = 1 # noqa: F841 result = pd.eval("x + 1", engine=engine, parser=parser) assert result == 2 df = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) df2 = pd.DataFrame(np.random.default_rng(2).standard_normal((5, 3))) # don't have the pandas parser msg = r"The '@' prefix is only supported by the pandas parser" with pytest.raises(SyntaxError, match=msg): df.query("(@df>0) & (@df2>0)", engine=engine, parser=parser) with pytest.raises(UndefinedVariableError, match="name 'df' is not defined"): df.query("(df>0) & (df2>0)", engine=engine, parser=parser) expected = df[(df > 0) & (df2 > 0)] result = pd.eval("df[(df > 0) & (df2 > 0)]", engine=engine, parser=parser) tm.assert_frame_equal(expected, result) expected = df[(df > 0) & (df2 > 0) & (df[df > 0] > 0)] result = pd.eval( "df[(df > 0) & (df2 > 0) & (df[df > 0] > 0)]", engine=engine, parser=parser ) tm.assert_frame_equal(expected, result) @skip_if_no_numexpr def test_query_numexpr_with_min_and_max_columns(self): df = pd.DataFrame({"min": [1, 2, 3], "max": [4, 5, 6]}) regex_to_match = ( r"Variables in expression \"\(min\) == \(1\)\" " r"overlap with builtins: \('min'\)" ) with pytest.raises(NumExprClobberingError, match=regex_to_match): df.query("min == 1") regex_to_match = ( r"Variables in expression \"\(max\) == \(1\)\" " r"overlap with builtins: \('max'\)" ) with pytest.raises(NumExprClobberingError, match=regex_to_match): df.query("max == 1") class TestDataFrameQueryPythonPandas(TestDataFrameQueryNumExprPandas): @pytest.fixture def engine(self): return "python" @pytest.fixture def parser(self): return "pandas" def test_query_builtin(self, engine, parser): n = m = 10 df = pd.DataFrame( np.random.default_rng(2).integers(m, size=(n, 3)), columns=list("abc") ) df.index.name = "sin" expected = df[df.index > 5] result = df.query("sin > 5", engine=engine, parser=parser) tm.assert_frame_equal(expected, result) class TestDataFrameQueryPythonPython(TestDataFrameQueryNumExprPython): @pytest.fixture def engine(self): return "python" @pytest.fixture def parser(self): return "python" def test_query_builtin(self, engine, parser): n = m = 10 df = pd.DataFrame( np.random.default_rng(2).integers(m, size=(n, 3)), columns=list("abc") ) df.index.name = "sin" expected = df[df.index > 5] result = df.query("sin > 5", engine=engine, parser=parser) tm.assert_frame_equal(expected, result) class TestDataFrameQueryStrings: def test_query_ordered_categorical_comparison(self, parser, engine): # GH#15186 comparing an ordered categorical column against a scalar # string in .query should work for all engines/parsers bands = [f"{i}AM" for i in range(10)] df = pd.DataFrame({"Time": bands}) df["Time"] = df["Time"].astype(pd.CategoricalDtype(bands, ordered=True)) result = df.query("Time >= '5AM'", parser=parser, engine=engine) expected = df[df["Time"] >= "5AM"] tm.assert_frame_equal(result, expected) def test_str_query_method(self, parser, engine): df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 1)), columns=["b"] ) df["strings"] = pd.Series(list("aabbccddee")) expect = df[df.strings == "a"] if parser != "pandas": col = "strings" lst = '"a"' lhs = [col] * 2 + [lst] * 2 rhs = lhs[::-1] eq, ne = "==", "!=" ops = 2 * ([eq, ne]) msg = r"'(Not)?In' nodes are not implemented" for lh, op_, rh in zip(lhs, ops, rhs, strict=True): ex = f"{lh} {op_} {rh}" with pytest.raises(NotImplementedError, match=msg): df.query( ex, engine=engine, parser=parser, local_dict={"strings": df.strings}, ) else: res = df.query('"a" == strings', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) res = df.query('strings == "a"', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) tm.assert_frame_equal(res, df[df.strings.isin(["a"])]) expect = df[df.strings != "a"] res = df.query('strings != "a"', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) res = df.query('"a" != strings', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) tm.assert_frame_equal(res, df[~df.strings.isin(["a"])]) def test_str_list_query_method(self, parser, engine): df = pd.DataFrame( np.random.default_rng(2).standard_normal((10, 1)), columns=["b"] ) df["strings"] = pd.Series(list("aabbccddee")) expect = df[df.strings.isin(["a", "b"])] if parser != "pandas": col = "strings" lst = '["a", "b"]' lhs = [col] * 2 + [lst] * 2 rhs = lhs[::-1] eq, ne = "==", "!=" ops = 2 * ([eq, ne]) msg = r"'(Not)?In' nodes are not implemented" for lh, ops_, rh in zip(lhs, ops, rhs, strict=True): ex = f"{lh} {ops_} {rh}" with pytest.raises(NotImplementedError, match=msg): df.query(ex, engine=engine, parser=parser) else: res = df.query('strings == ["a", "b"]', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) res = df.query('["a", "b"] == strings', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) expect = df[~df.strings.isin(["a", "b"])] res = df.query('strings != ["a", "b"]', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) res = df.query('["a", "b"] != strings', engine=engine, parser=parser) tm.assert_frame_equal(res, expect) def test_query_str_accessor_method(self, parser, engine): # GH#26549 df = pd.DataFrame({"colA": list("0123456789")}) result = df.query('colA.str.contains("0")', engine=engine, parser=parser) tm.assert_frame_equal(result, df[df["colA"].str.contains("0")]) def test_query_str_contains_word_boundary(self, parser, engine): # GH#32589 was a non-raw outer string, not a pandas bug; this pins that # _preparse leaves a backslash escape inside a string literal alone df = pd.DataFrame({"A": ["A test case", "Another Testing Case"]}) pat = r"\btest\b" result = df.query( f"A.str.contains(r'{pat}', regex=True, case=False)", engine=engine, parser=parser, ) expected = df[df["A"].str.contains(pat, regex=True, case=False)] tm.assert_frame_equal(result, expected) def test_query_str_slice_negative_index(self, parser, engine): # GH#49905 visit_Call's unary argument, as in # test_eval.py::test_unary_in_function; the a.str[-1:] spelling in the # same issue is test_eval.py::test_slice_subscript_with_unary_bound df = pd.DataFrame({"a": ["example", "zzz"]}) result = df.query( 'a.str.slice(0, -1).str.contains("e")', engine=engine, parser=parser ) expected = df[df["a"].str.slice(0, -1).str.contains("e")] tm.assert_frame_equal(result, expected) def test_query_with_string_columns(self, parser, engine): df = pd.DataFrame( { "a": list("aaaabbbbcccc"), "b": list("aabbccddeeff"), "c": np.random.default_rng(2).integers(5, size=12), "d": np.random.default_rng(2).integers(9, size=12), } ) if parser == "pandas": res = df.query("a in b", parser=parser, engine=engine) expec = df[df.a.isin(df.b)] tm.assert_frame_equal(res, expec) res = df.query("a in b and c < d", parser=parser, engine=engine) expec = df[df.a.isin(df.b) & (df.c < df.d)] tm.assert_frame_equal(res, expec) else: msg = r"'(Not)?In' nodes are not implemented" with pytest.raises(NotImplementedError, match=msg): df.query("a in b", parser=parser, engine=engine) msg = r"'BoolOp' nodes are not implemented" with pytest.raises(NotImplementedError, match=msg): df.query("a in b and c < d", parser=parser, engine=engine) def test_object_array_eq_ne(self, parser, engine): df = pd.DataFrame( { "a": list("aaaabbbbcccc"), "b": list("aabbccddeeff"), "c": np.random.default_rng(2).integers(5, size=12), "d": np.random.default_rng(2).integers(9, size=12), } ) res = df.query("a == b", parser=parser, engine=engine) exp = df[df.a == df.b] tm.assert_frame_equal(res, exp) res = df.query("a != b", parser=parser, engine=engine) exp = df[df.a != df.b] tm.assert_frame_equal(res, exp) def test_query_with_nested_strings(self, parser, engine): skip_if_no_pandas_parser(parser) events = [ f"page {n} {act}" for n in range(1, 4) for act in ["load", "exit"] ] * 2 stamps1 = pd.date_range("2014-01-01 0:00:01", freq="30s", periods=6) stamps2 = pd.date_range("2014-02-01 1:00:01", freq="30s", periods=6) df = pd.DataFrame( { "id": np.arange(1, 7).repeat(2), "event": events, "timestamp": stamps1.append(stamps2), } ) expected = df[df.event == '"page 1 load"'] res = df.query("""'"page 1 load"' in event""", parser=parser, engine=engine) tm.assert_frame_equal(expected, res) def test_query_with_nested_special_character(self, parser, engine): skip_if_no_pandas_parser(parser) df = pd.DataFrame({"a": ["a", "b", "test & test"], "b": [1, 2, 3]}) res = df.query('a == "test & test"', parser=parser, engine=engine) expec = df[df.a == "test & test"] tm.assert_frame_equal(res, expec) @pytest.mark.parametrize( "op, func", [ ["<", operator.lt], [">", operator.gt], ["<=", operator.le], [">=", operator.ge], ], ) def test_query_lex_compare_strings(self, parser, engine, op, func): a = pd.Series(np.random.default_rng(2).choice(list("abcde"), 20)) b = pd.Series(np.arange(a.size)) df = pd.DataFrame({"X": a, "Y": b}) res = df.query(f'X {op} "d"', engine=engine, parser=parser) expected = df[func(df.X, "d")] tm.assert_frame_equal(res, expected) def test_query_single_element_booleans(self, parser, engine): columns = "bid", "bidsize", "ask", "asksize" data = np.random.default_rng(2).integers(2, size=(1, len(columns))).astype(bool) df = pd.DataFrame(data, columns=columns) res = df.query("bid & ask", engine=engine, parser=parser) expected = df[df.bid & df.ask] tm.assert_frame_equal(res, expected) def test_query_string_scalar_variable(self, parser, engine): skip_if_no_pandas_parser(parser) df = pd.DataFrame( { "Symbol": ["BUD US", "BUD US", "IBM US", "IBM US"], "Price": [109.70, 109.72, 183.30, 183.35], } ) e = df[df.Symbol == "BUD US"] symb = "BUD US" # noqa: F841 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): # GITHUB ISSUE #31516 parser = "pandas" engine = "python" expected = {i: value for i, value in enumerate(in_list) if value == "asdf"} df_expected = pd.DataFrame({"a": expected}, dtype="string") df_expected.index = df_expected.index.astype("int64") df = pd.DataFrame({"a": in_list}, dtype="string") df.index = pd.Index(list(df.index), dtype=df.index.dtype) 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 def frame(self): return pd.DataFrame( np.random.default_rng(2).standard_normal((10, 3)), columns=list("abc") ) def test_simple_expr(self, frame, parser, engine): res = frame.eval("a + b", engine=engine, parser=parser) expect = frame.a + frame.b tm.assert_series_equal(res, expect) def test_bool_arith_expr(self, frame, parser, engine): res = frame.eval("a[a < 1] + b", engine=engine, parser=parser) expect = frame.a[frame.a < 1] + frame.b tm.assert_series_equal(res, expect) @pytest.mark.parametrize("op", ["+", "-", "*", "/"]) def test_invalid_type_for_operator_raises(self, parser, engine, op): df = pd.DataFrame({"a": [1, 2], "b": ["c", "d"]}) msg = r"unsupported operand type\(s\) for .+: '.+' and '.+'" with pytest.raises(TypeError, match=msg): df.eval(f"a {op} b", engine=engine, parser=parser) class TestDataFrameQueryBacktickQuoting: @pytest.fixture def df(self): """ Yields a dataframe with strings that may or may not need escaping by backticks. The last two columns cannot be escaped by backticks and should raise a ValueError. """ return pd.DataFrame( { "A": [1, 2, 3], "B B": [3, 2, 1], "C C": [4, 5, 6], "C C": [7, 4, 3], "C_C": [8, 9, 10], "D_D D": [11, 1, 101], "E.E": [6, 3, 5], "F-F": [8, 1, 10], "1e1": [2, 4, 8], "def": [10, 11, 2], "A (x)": [4, 1, 3], "B(x)": [1, 1, 5], "B (x)": [2, 7, 4], " &^ :!€$?(} > <++*'' ": [2, 5, 6], "": [10, 11, 1], " A": [4, 7, 9], " ": [1, 2, 1], "it's": [6, 3, 1], "that's": [9, 1, 8], "☺": [8, 7, 6], "xy （z）": [1, 2, 3], # noqa: RUF001 "xy （z\\uff09": [4, 5, 6], # noqa: RUF001 "foo#bar": [2, 4, 5], 1: [5, 7, 9], } ) def test_single_backtick_variable_query(self, df): res = df.query("1 < `B B`") expect = df[1 < df["B B"]] tm.assert_frame_equal(res, expect) def test_two_backtick_variables_query(self, df): res = df.query("1 < `B B` and 4 < `C C`") expect = df[(1 < df["B B"]) & (4 < df["C C"])] tm.assert_frame_equal(res, expect) def test_single_backtick_variable_expr(self, df): res = df.eval("A + `B B`") expect = df["A"] + df["B B"] tm.assert_series_equal(res, expect) def test_two_backtick_variables_expr(self, df): res = df.eval("`B B` + `C C`") expect = df["B B"] + df["C C"] tm.assert_series_equal(res, expect) def test_already_underscore_variable(self, df): res = df.eval("`C_C` + A") expect = df["C_C"] + df["A"] tm.assert_series_equal(res, expect) def test_same_name_but_underscores(self, df): res = df.eval("C_C + `C C`") expect = df["C_C"] + df["C C"] tm.assert_series_equal(res, expect) def test_mixed_underscores_and_spaces(self, df): res = df.eval("A + `D_D D`") expect = df["A"] + df["D_D D"] tm.assert_series_equal(res, expect) def test_backtick_quote_name_with_no_spaces(self, df): res = df.eval("A + `C_C`") expect = df["A"] + df["C_C"] tm.assert_series_equal(res, expect) def test_special_characters(self, df): res = df.eval("`E.E` + `F-F` - A") expect = df["E.E"] + df["F-F"] - df["A"] tm.assert_series_equal(res, expect) def test_start_with_digit(self, df): res = df.eval("A + `1e1`") expect = df["A"] + df["1e1"] tm.assert_series_equal(res, expect) def test_keyword(self, df): res = df.eval("A + `def`") expect = df["A"] + df["def"] tm.assert_series_equal(res, expect) def test_unneeded_quoting(self, df): res = df.query("`A` > 2") expect = df[df["A"] > 2] tm.assert_frame_equal(res, expect) def test_parenthesis(self, df): res = df.query("`A (x)` > 2") expect = df[df["A (x)"] > 2] tm.assert_frame_equal(res, expect) def test_empty_string(self, df): res = df.query("`` > 5") expect = df[df[""] > 5] tm.assert_frame_equal(res, expect) def test_multiple_spaces(self, df): res = df.query("`C C` > 5") expect = df[df["C C"] > 5] tm.assert_frame_equal(res, expect) def test_start_with_spaces(self, df): res = df.eval("` A` + ` `") expect = df[" A"] + df[" "] tm.assert_series_equal(res, expect) def test_ints(self, df): res = df.query("`1` == 7") expect = df[df[1] == 7] tm.assert_frame_equal(res, expect) def test_lots_of_operators_string(self, df): res = df.query("` &^ :!€$?(} > <++*'' ` > 4") expect = df[df[" &^ :!€$?(} > <++*'' "] > 4] tm.assert_frame_equal(res, expect) def test_missing_attribute(self, df): message = "module 'pandas' has no attribute 'thing'" with pytest.raises(AttributeError, match=message): df.eval("@pd.thing") def test_quote(self, df): res = df.query("`it's` > `that's`") expect = df[df["it's"] > df["that's"]] tm.assert_frame_equal(res, expect) def test_character_outside_range_smiley(self, df): res = df.query("`☺` > 4") expect = df[df["☺"] > 4] tm.assert_frame_equal(res, expect) def test_character_outside_range_2_byte_parens(self, df): # GH 49633 res = df.query("`xy （z）` == 2") # noqa: RUF001 expect = df[df["xy （z）"] == 2] # noqa: RUF001 tm.assert_frame_equal(res, expect) def test_character_outside_range_and_actual_backslash(self, df): # GH 49633 res = df.query("`xy （z\\uff09` == 2") # noqa: RUF001 expect = df[df["xy \uff08z\\uff09"] == 2] tm.assert_frame_equal(res, expect) def test_hashtag(self, df): res = df.query("`foo#bar` > 4") expect = df[df["foo#bar"] > 4] tm.assert_frame_equal(res, expect) def test_expr_with_column_name_with_hashtag_character(self): # GH 59285 df = pd.DataFrame((1, 2, 3), columns=["a#"]) result = df.query("`a#` < 2") expected = df[df["a#"] < 2] tm.assert_frame_equal(result, expected) def test_expr_with_comment(self): # GH 59285 df = pd.DataFrame((1, 2, 3), columns=["a#"]) result = df.query("`a#` < 2 # This is a comment") expected = df[df["a#"] < 2] tm.assert_frame_equal(result, expected) def test_expr_with_column_name_with_backtick_and_hash(self): # GH 59285 df = pd.DataFrame((1, 2, 3), columns=["a`#b"]) result = df.query("`a``#b` < 2") expected = df[df["a`#b"] < 2] tm.assert_frame_equal(result, expected) def test_expr_with_column_name_with_backtick(self): # GH 59285 df = pd.DataFrame({"a`b": (1, 2, 3), "ab": (4, 5, 6)}) result = df.query("`a``b` < 2") # Note: Formatting checks may wrongly consider the above ``inline code``. expected = df[df["a`b"] < 2] tm.assert_frame_equal(result, expected) def test_expr_with_string_with_backticks(self): # GH 59285 df = pd.DataFrame(("`", "`````", "``````````"), columns=["#backticks"]) result = df.query("'```' < `#backticks`") expected = df["```" < df["#backticks"]] tm.assert_frame_equal(result, expected) def test_expr_with_string_with_backticked_substring_same_as_column_name(self): # GH 59285 df = pd.DataFrame(("`", "`````", "``````````"), columns=["#backticks"]) result = df.query("'`#backticks`' < `#backticks`") expected = df["`#backticks`" < df["#backticks"]] tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "col1,col2,expr", [ ("it's", "that's", "`it's` < `that's`"), ('it"s', 'that"s', '`it"s` < `that"s`'), ("it's", 'that\'s "nice"', "`it's` < `that's \"nice\"`"), ("it's", "that's #cool", "`it's` < `that's #cool` # This is a comment"), ], ) def test_expr_with_column_names_with_special_characters(self, col1, col2, expr): # GH 59285 df = pd.DataFrame( [ {col1: 1, col2: 2}, {col1: 3, col2: 4}, {col1: -1, col2: -2}, {col1: -3, col2: -4}, ] ) result = df.query(expr) expected = df[df[col1] < df[col2]] tm.assert_frame_equal(result, expected) def test_expr_with_no_backticks(self): # GH 59285 df = pd.DataFrame(("aaa", "vvv", "zzz"), columns=["column_name"]) result = df.query("'value' < column_name") expected = df["value" < df["column_name"]] tm.assert_frame_equal(result, expected) def test_expr_with_no_quotes_and_backtick_is_unmatched(self): # GH 59285 df = pd.DataFrame((1, 5, 10), columns=["column-name"]) with pytest.raises((SyntaxError, TokenError), match="invalid syntax"): df.query("5 < `column-name") def test_expr_with_no_quotes_and_backtick_is_matched(self): # GH 59285 df = pd.DataFrame((1, 5, 10), columns=["column-name"]) result = df.query("5 < `column-name`") expected = df[5 < df["column-name"]] tm.assert_frame_equal(result, expected) def test_expr_with_backtick_opened_before_quote_and_backtick_is_unmatched(self): # GH 59285 df = pd.DataFrame((1, 5, 10), columns=["It's"]) with pytest.raises( (SyntaxError, TokenError), match="unterminated string literal" ): df.query("5 < `It's") def test_expr_with_backtick_opened_before_quote_and_backtick_is_matched(self): # GH 59285 df = pd.DataFrame((1, 5, 10), columns=["It's"]) result = df.query("5 < `It's`") expected = df[5 < df["It's"]] tm.assert_frame_equal(result, expected) def test_expr_with_quote_opened_before_backtick_and_quote_is_unmatched(self): # GH 59285 df = pd.DataFrame(("aaa", "vvv", "zzz"), columns=["column-name"]) with pytest.raises( (SyntaxError, TokenError), match="unterminated string literal" ): df.query("`column-name` < 'It`s that\\'s \"quote\" #hash") def test_expr_with_quote_opened_before_backtick_and_quote_is_matched_at_end(self): # GH 59285 df = pd.DataFrame(("aaa", "vvv", "zzz"), columns=["column-name"]) result = df.query("`column-name` < 'It`s that\\'s \"quote\" #hash'") expected = df[df["column-name"] < 'It`s that\'s "quote" #hash'] tm.assert_frame_equal(result, expected) def test_expr_with_quote_opened_before_backtick_and_quote_is_matched_in_mid(self): # GH 59285 df = pd.DataFrame(("aaa", "vvv", "zzz"), columns=["column-name"]) result = df.query("'It`s that\\'s \"quote\" #hash' < `column-name`") expected = df['It`s that\'s "quote" #hash' < df["column-name"]] tm.assert_frame_equal(result, expected) def test_call_non_named_expression(self, df): """ Only attributes and variables ('named functions') can be called. .__call__() is not an allowed attribute because that would allow calling anything. https://github.com/pandas-dev/pandas/pull/32460 """ def func(*_): return 1 funcs = [func] # noqa: F841 df.eval("@func()") with pytest.raises(TypeError, match="Only named functions are supported"): df.eval("@funcs[0]()") with pytest.raises(TypeError, match="Only named functions are supported"): df.eval("@funcs[0].__call__()") def test_ea_dtypes(self, any_numeric_ea_and_arrow_dtype): # GH#29618 df = pd.DataFrame( [[1, 2], [3, 4]], columns=["a", "b"], dtype=any_numeric_ea_and_arrow_dtype ) warning = RuntimeWarning if NUMEXPR_INSTALLED else None msg = ( "Engine has switched to 'python' because numexpr does not " "support extension array dtypes. Please set your engine " "to python manually." ) with tm.assert_produces_warning(warning, match=msg): result = df.eval("c = b - a") expected = pd.DataFrame( [[1, 2, 1], [3, 4, 1]], columns=["a", "b", "c"], dtype=any_numeric_ea_and_arrow_dtype, ) tm.assert_frame_equal(result, expected) def test_ea_dtypes_and_scalar(self): # GH#29618 df = pd.DataFrame([[1, 2], [3, 4]], columns=["a", "b"], dtype="Float64") warning = RuntimeWarning if NUMEXPR_INSTALLED else None msg = ( "Engine has switched to 'python' because numexpr does not " "support extension array dtypes. Please set your engine " "to python manually." ) with tm.assert_produces_warning(warning, match=msg): result = df.eval("c = b - 1") expected = pd.DataFrame( [[1, 2, 1], [3, 4, 3]], columns=["a", "b", "c"], dtype="Float64" ) tm.assert_frame_equal(result, expected) def test_ea_dtypes_and_scalar_operation(self, any_numeric_ea_and_arrow_dtype): # GH#29618 df = pd.DataFrame( [[1, 2], [3, 4]], columns=["a", "b"], dtype=any_numeric_ea_and_arrow_dtype ) result = df.eval("c = 2 - 1") expected = pd.DataFrame( { "a": pd.Series([1, 3], dtype=any_numeric_ea_and_arrow_dtype), "b": pd.Series([2, 4], dtype=any_numeric_ea_and_arrow_dtype), "c": pd.Series([1, 1], dtype=result["c"].dtype), } ) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("dtype", ["int64", "Int64", "int64[pyarrow]"]) def test_query_ea_dtypes(self, dtype): if dtype == "int64[pyarrow]": pytest.importorskip("pyarrow") # GH#50261 df = pd.DataFrame({"a": [1, 2]}, dtype=dtype) ref = {2} # noqa: F841 warning = RuntimeWarning if dtype == "Int64" and NUMEXPR_INSTALLED else None msg = ( "Engine has switched to 'python' because numexpr does not " "support extension array dtypes. Please set your engine " "to python manually." ) with tm.assert_produces_warning(warning, match=msg): result = df.query("a in @ref") expected = pd.DataFrame({"a": [2]}, index=range(1, 2), dtype=dtype) tm.assert_frame_equal(result, expected) def test_query_ea_dtype_scalar_comparison(self, any_numeric_ea_and_arrow_dtype): # GH#25369 df = pd.DataFrame( {"col_test": [4, None, 6]}, dtype=any_numeric_ea_and_arrow_dtype ) warning = RuntimeWarning if NUMEXPR_INSTALLED else None msg = ( "Engine has switched to 'python' because numexpr does not " "support extension array dtypes. Please set your engine " "to python manually." ) with tm.assert_produces_warning(warning, match=msg): result = df.query("col_test != 6") tm.assert_frame_equal(result, df.iloc[:1]) @pytest.mark.parametrize("engine", ["python", "numexpr"]) @pytest.mark.parametrize("dtype", ["int64", "Int64", "int64[pyarrow]"]) def test_query_ea_equality_comparison(self, dtype, engine): # GH#50261 warning = RuntimeWarning if engine == "numexpr" else None if engine == "numexpr" and not NUMEXPR_INSTALLED: pytest.skip("numexpr not installed") if dtype == "int64[pyarrow]": pytest.importorskip("pyarrow") df = pd.DataFrame( { "A": pd.Series([1, 1, 2], dtype="Int64"), "B": pd.Series([1, 2, 2], dtype=dtype), } ) msg = ( "Engine has switched to 'python' because numexpr does not " "support extension array dtypes. Please set your engine " "to python manually." ) with tm.assert_produces_warning(warning, match=msg): result = df.query("A == B", engine=engine) expected = pd.DataFrame( { "A": pd.Series([1, 2], dtype="Int64", index=range(0, 4, 2)), "B": pd.Series([1, 2], dtype=dtype, index=range(0, 4, 2)), } ) tm.assert_frame_equal(result, expected) def test_all_nat_in_object(self): # GH#57068 now = pd.Timestamp("2011-01-01", tz="UTC") # noqa: F841 df = pd.DataFrame({"a": pd.to_datetime([None, None], utc=True)}, dtype=object) result = df.query("a > @now") expected = pd.DataFrame({"a": []}, dtype=object) tm.assert_frame_equal(result, expected) class TestDataFrameQueryInWithColumnRefs: # GH#65357 - df.query('col in (col1, col2)') with column references. @pytest.fixture(autouse=True) def skip_if_no_pandas_parser(self, parser): if parser != "pandas": pytest.skip(f"cannot evaluate with parser={parser}") def test_in_column_refs_numeric_first_row(self, engine, parser): df = pd.DataFrame({"a": [1, 2, 3], "b": [1, 4, 9], "c": [1, 8, 27]}) result = df.query("b in (a, c)", engine=engine, parser=parser) expected = df[(df["b"] == df["a"]) | (df["b"] == df["c"])] tm.assert_frame_equal(result, expected) def test_in_column_refs_numeric_all_rows(self, engine, parser): df = pd.DataFrame({"a": [1, 2, 3], "b": [1, 2, 3], "c": [10, 20, 30]}) result = df.query("b in (a, c)", engine=engine, parser=parser) expected = df tm.assert_frame_equal(result, expected) def test_in_column_refs_numeric_no_match(self, engine, parser): df = pd.DataFrame({"a": [1, 2, 3], "b": [5, 6, 7], "c": [10, 20, 30]}) result = df.query("b in (a, c)", engine=engine, parser=parser) expected = df.iloc[:0] tm.assert_frame_equal(result, expected) def test_in_column_refs_string(self, engine, parser): df = pd.DataFrame( {"a": ["a", "a", "b"], "b": ["a", "b", "c"], "c": ["b", "b", "c"]} ) result = df.query("b in (a, c)", engine=engine, parser=parser) expected = df[(df["b"] == df["a"]) | (df["b"] == df["c"])] tm.assert_frame_equal(result, expected) @td.skip_if_no("pyarrow") def test_in_column_refs_string_pyarrow(self, engine, parser): df = pd.DataFrame( {"a": ["a", "a", "b"], "b": ["a", "b", "c"], "c": ["b", "b", "c"]}, dtype="string[pyarrow]", ) warning = RuntimeWarning if engine == "numexpr" else None msg = ( "Engine has switched to 'python' because numexpr does not " "support extension array dtypes. Please set your engine " "to python manually." ) with tm.assert_produces_warning(warning, match=msg): result = df.query("b in (a, c)", engine=engine, parser=parser) expected = df[(df["b"] == df["a"]) | (df["b"] == df["c"])] tm.assert_frame_equal(result, expected) def test_not_in_column_refs_numeric(self, engine, parser): df = pd.DataFrame({"a": [1, 2, 3], "b": [1, 4, 9], "c": [1, 8, 27]}) result = df.query("b not in (a, c)", engine=engine, parser=parser) expected = df[(df["b"] != df["a"]) & (df["b"] != df["c"])] tm.assert_frame_equal(result, expected) def test_not_in_column_refs_string(self, engine, parser): df = pd.DataFrame( {"a": ["x", "y", "z"], "b": ["a", "y", "z"], "c": ["b", "b", "c"]} ) result = df.query("b not in (a, c)", engine=engine, parser=parser) expected = df[(df["b"] != df["a"]) & (df["b"] != df["c"])] tm.assert_frame_equal(result, expected) def test_in_column_ref_and_literal(self, engine, parser): df = pd.DataFrame({"a": [1, 2, 3], "b": [1, 99, 5]}) result = df.query("b in (a, 99)", engine=engine, parser=parser) expected = df[(df["b"] == df["a"]) | (df["b"] == 99)] tm.assert_frame_equal(result, expected) def test_in_pure_literals(self, engine, parser): df = pd.DataFrame({"a": [1, 2, 3, 4]}) result = df.query("a in (2, 3)", engine=engine, parser=parser) expected = df[df["a"].isin([2, 3])] tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "dtype, values", [ ("datetime64[ns]", ["2020-01-01", "2020-01-02"]), ("timedelta64[ns]", ["5 days", "3 days"]), ("period[D]", ["2020-01-01", "2020-01-02"]), ], ) @pytest.mark.xfail( reason="GH#23420, GH#35595 _rewrite_membership_op turns == into isin, " "which does not parse the string" ) def test_query_datetimelike_eq_string(dtype, values, engine, parser): skip_if_no_pandas_parser(parser) df = pd.DataFrame({"a": pd.array(values, dtype=dtype)}) result = df.query(f'a == "{values[0]}"', engine=engine, parser=parser) tm.assert_frame_equal(result, df.iloc[[0]]) @pytest.mark.parametrize( "op, func", [("<", operator.lt), ("<=", operator.le), (">", operator.gt), (">=", operator.ge)], ) def test_query_timedelta_compared_to_string(op, func, engine, parser): # GH#23420 the ordered comparisons raised ValueError on timedelta; # == is still wrong, see test_query_datetimelike_eq_string df = pd.DataFrame({"dt": pd.to_timedelta(["5 days", "3 days"])}) result = df.query(f'dt {op} "4 days"', engine=engine, parser=parser) expected = df[func(df["dt"], pd.Timedelta("4 days"))] tm.assert_frame_equal(result, expected)