Repository: dask/dask

Issue:
Improve `DataFrame.query` docstring with local variable example
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
Prepare Dask distributed:
```
from dask.distributed import Client

client = Client(n_workers=16,
                threads_per_worker=1)
```

This works:
```python
import pandas as pd

data = {'col_1': [3, 2, 1, 0], 'col_2': ['a', 'b', 'c', 'd']}
df = pd.DataFrame.from_dict(data)
pick = 'c'
df.query('col_2 == @pick')
```
This doesn't work:
```python
dd_temp = dd.from_pandas(df, npartitions=2)
dd_temp.query('col_2 == @pick')
```
Throws UndefinedVariableError("local variable 'pick' is not defined").  

**What you expected to happen**:
Dask version of DataFrame.query should work as in pandas.

**Minimal Complete Verifiable Example**:
See above

**Anything else we need to know?**:
No

**Environment**:

- Dask version: 2021.8.0
- Python version: 3.9.6
- Operating System: Linux (Manjaro)
- Install method (conda, pip, source): conda-forge



Discussion:
Hmm yeah I don't think the `@` syntax is necessarily supported in dask. As a workaround either of these will work:

```python
dd_temp.query(f'col_2 == "{pick}"').compute()
```

```python
dd_temp[dd_temp["col_2"] == pick].compute()
```
Or you can do 

```python
dd_temp.query('col_2 == @pick', local_dict={"pick": pick}).compute()
> Hmm yeah I don't think the `@` syntax is necessarily supported in dask. As a workaround either of these will work:
> 
> ```python
> dd_temp.query(f'col_2 == "{pick}"').compute()
> ```
> 
> ```python
> dd_temp[dd_temp["col_2"] == pick].compute()
> ```

Sometimes there will be complex queries. Besides, f-string does not always work, for example, if you want to use datetime object into the query. 
Will @jsignell's other suggestion (xref https://github.com/dask/dask/issues/8156#issuecomment-921153847) work in that case?

```python
dd_temp.query('col_2 == @pick', local_dict={"pick": pick}).compute()
```
 @umonaca Checking in here, Do you think the second suggestion made here https://github.com/dask/dask/issues/8156#issuecomment-921153847 would work in that case?  
I'll look into it tomorrow.
Just checking in, @umonaca did this work out? 
> Just checking in, @umonaca did this work out?

The solution works. However, I think you should document this parameter in `dask.dataframe.DataFrame.query`.
Ohh! That's true https://docs.dask.org/en/latest/generated/dask.dataframe.DataFrame.query.html is very thin! I only realized that pattern was supported because I was looking around in the tests. Please feel free to open a pull request with an example of that pattern @umonaca!

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -4814,6 +4814,17 @@ def query(self, expr, **kwargs):
 
         Blocked version of pd.DataFrame.query
 
+        Parameters
+        ----------
+        expr: str
+            The query string to evaluate.
+            You can refer to column names that are not valid Python variable names
+            by surrounding them in backticks.
+            Dask does not fully support referring to variables using the '@' character,
+            use f-strings or the ``local_dict`` keyword argument instead.
+
+        Notes
+        -----
         This is like the sequential version except that this will also happen
         in many threads.  This may conflict with ``numexpr`` which will use
         multiple threads itself.  We recommend that you set ``numexpr`` to use a
@@ -4827,6 +4838,46 @@ def query(self, expr, **kwargs):
         See also
         --------
         pandas.DataFrame.query
+        pandas.eval
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import dask.dataframe as dd
+        >>> df = pd.DataFrame({'x': [1, 2, 1, 2],
+        ...                    'y': [1, 2, 3, 4],
+        ...                    'z z': [4, 3, 2, 1]})
+        >>> ddf = dd.from_pandas(df, npartitions=2)
+
+        Refer to column names directly:
+
+        >>> ddf.query('y > x').compute()
+           x  y  z z
+        2  1  3    2
+        3  2  4    1
+
+        Refer to column name using backticks:
+
+        >>> ddf.query('`z z` > x').compute()
+           x  y  z z
+        0  1  1    4
+        1  2  2    3
+        2  1  3    2
+
+        Refer to variable name using f-strings:
+
+        >>> value = 1
+        >>> ddf.query(f'x == {value}').compute()
+           x  y  z z
+        0  1  1    4
+        2  1  3    2
+
+        Refer to variable name using ``local_dict``:
+
+        >>> ddf.query('x == @value', local_dict={"value": value}).compute()
+           x  y  z z
+        0  1  1    4
+        2  1  3    2
         """
         return self.map_partitions(M.query, expr, **kwargs)
 


Test patch: