Skip to content
Merged
96 changes: 86 additions & 10 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -4481,18 +4481,58 @@ def _get_item(self, item: Hashable) -> Series:

@overload
def query(
self, expr: str, *, inplace: Literal[False] = ..., **kwargs
self,
expr: str,
*,
parser: Literal["pandas", "python"] = ...,
engine: Literal["python", "numexpr"] | None = ...,
local_dict: dict[str, Any] | None = ...,
global_dict: dict[str, Any] | None = ...,
resolvers: list[Mapping] | None = ...,
level: int = ...,
inplace: Literal[False] = ...,
) -> DataFrame: ...

@overload
def query(self, expr: str, *, inplace: Literal[True], **kwargs) -> None: ...
def query(
self,
expr: str,
*,
parser: Literal["pandas", "python"] = ...,
engine: Literal["python", "numexpr"] | None = ...,
local_dict: dict[str, Any] | None = ...,
global_dict: dict[str, Any] | None = ...,
resolvers: list[Mapping] | None = ...,
level: int = ...,
inplace: Literal[True],
) -> None: ...

@overload
def query(
self, expr: str, *, inplace: bool = ..., **kwargs
self,
expr: str,
*,
parser: Literal["pandas", "python"] = ...,
engine: Literal["python", "numexpr"] | None = ...,
local_dict: dict[str, Any] | None = ...,
global_dict: dict[str, Any] | None = ...,
resolvers: list[Mapping] | None = ...,
level: int = ...,
inplace: bool = ...,
) -> DataFrame | None: ...

def query(self, expr: str, *, inplace: bool = False, **kwargs) -> DataFrame | None:
def query(
self,
expr: str,
*,
parser: Literal["pandas", "python"] = "pandas",
engine: Literal["python", "numexpr"] | None = None,
local_dict: dict[str, Any] | None = None,
global_dict: dict[str, Any] | None = None,
resolvers: list[Mapping] | None = None,
level: int = 0,
inplace: bool = False,
) -> DataFrame | None:
"""
Query the columns of a DataFrame with a boolean expression.

Expand All @@ -4511,11 +4551,41 @@ def query(self, expr: str, *, inplace: bool = False, **kwargs) -> DataFrame | No

See the documentation for :meth:`DataFrame.eval` for details on
referring to column names and variables in the query string.
parser : {'pandas', 'python'}, default 'pandas'
The parser to use to construct the syntax tree from the expression. The
default of ``'pandas'`` parses code slightly different than standard
Python. Alternatively, you can parse an expression using the
``'python'`` parser to retain strict Python semantics. See the
:ref:`enhancing performance <enhancingperf.eval>` documentation for
more details.
engine : {'python', 'numexpr'}, default 'numexpr'

The engine used to evaluate the expression. Supported engines are

- None : tries to use ``numexpr``, falls back to ``python``
- ``'numexpr'`` : This default engine evaluates pandas objects using
numexpr for large speed ups in complex expressions with large frames.
- ``'python'`` : Performs operations as if you had ``eval``'d in top
level python. This engine is generally not that useful.

More backends may be available in the future.
local_dict : dict or None, optional
A dictionary of local variables, taken from locals() by default.
global_dict : dict or None, optional
A dictionary of global variables, taken from globals() by default.
resolvers : list of dict-like or None, optional
A list of objects implementing the ``__getitem__`` special method that
you can use to inject an additional collection of namespaces to use for
variable lookup. For example, this is used in the
:meth:`~DataFrame.query` method to inject the
``DataFrame.index`` and ``DataFrame.columns``
variables that refer to their respective :class:`~pandas.DataFrame`
instance attributes.
level : int, optional
The number of prior stack frames to traverse and add to the current
scope. Most users will **not** need to change this parameter.
inplace : bool
Whether to modify the DataFrame rather than creating a new one.
**kwargs
See the documentation for :func:`eval` for complete details
on the keyword arguments accepted by :meth:`DataFrame.query`.

Returns
-------
Expand Down Expand Up @@ -4628,10 +4698,16 @@ def query(self, expr: str, *, inplace: bool = False, **kwargs) -> DataFrame | No
if not isinstance(expr, str):
msg = f"expr must be a string to be evaluated, {type(expr)} given"
raise ValueError(msg)
kwargs["level"] = kwargs.pop("level", 0) + 1
kwargs["target"] = None

res = self.eval(expr, **kwargs)
res = self.eval(
expr,
level=level + 1,
parser=parser,
engine=engine,
local_dict=local_dict,
global_dict=global_dict,
resolvers=resolvers or (),
)

try:
result = self.loc[res]
Expand Down
Loading