Skip to content

BUG: Solves errors when calling series methods in DataFrame.query with numexpr #43301

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Sep 25, 2021
Merged
1 change: 1 addition & 0 deletions doc/source/whatsnew/v1.4.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -426,6 +426,7 @@ Indexing
- Bug in :meth:`Index.get_indexer_non_unique` when index contains multiple ``np.nan`` (:issue:`35392`)
- Bug in :meth:`DataFrame.query` did not handle the degree sign in a backticked column name, such as \`Temp(°C)\`, used in an expression to query a dataframe (:issue:`42826`)
- Bug in :meth:`DataFrame.drop` where the error message did not show missing labels with commas when raising ``KeyError`` (:issue:`42881`)
- Bug in :meth:`DataFrame.query` where method calls in query strings led to errors when the ``numexpr`` package was installed. (:issue:`22435`)
- Bug in :meth:`DataFrame.nlargest` and :meth:`Series.nlargest` where sorted result did not count indexes containing ``np.nan`` (:issue:`28984`)


Expand Down
3 changes: 2 additions & 1 deletion pandas/core/computation/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -702,7 +702,8 @@ def visit_Call(self, node, side=None, **kwargs):
if key.arg:
kwargs[key.arg] = self.visit(key.value).value

return self.const_type(res(*new_args, **kwargs), self.env)
name = self.env.add_tmp(res(*new_args, **kwargs))
return self.term_type(name=name, env=self.env)

def translate_In(self, op):
return op
Expand Down
20 changes: 20 additions & 0 deletions pandas/tests/frame/test_query_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -731,6 +731,26 @@ def test_check_tz_aware_index_query(self, tz_aware_fixture):
result = df.reset_index().query('"2018-01-03 00:00:00+00" < time')
tm.assert_frame_equal(result, expected)

def test_method_calls_in_query(self):
# https://github.com/pandas-dev/pandas/issues/22435
n = 10
df = DataFrame({"a": 2 * np.random.rand(n), "b": np.random.rand(n)})
expected = df[df["a"].astype("int") == 0]
result = df.query(
"a.astype('int') == 0", engine=self.engine, parser=self.parser
)
tm.assert_frame_equal(result, expected)

df = DataFrame(
{
"a": np.where(np.random.rand(n) < 0.5, np.nan, np.random.randn(n)),
"b": np.random.randn(n),
}
)
expected = df[df["a"].notnull()]
result = df.query("a.notnull()", engine=self.engine, parser=self.parser)
tm.assert_frame_equal(result, expected)


@td.skip_if_no_ne
class TestDataFrameQueryNumExprPython(TestDataFrameQueryNumExprPandas):
Expand Down