NVIDIA / NVIDIA/cudf

[FEA] DataFrame.query support for math ops

Open
#12,510 0 comments 0 reactions 0 assignees View on GitHub
0 - Backlog feature request Python
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**
working with `import cudf as pd`

**Describe the solution you'd like**
support for math ops matching `pandas`

```
In [1]: import cudf as pd

In [2]: pd.__version__
Out[2]: '22.12.0'

In [3]: df = pd.DataFrame({'a': [1**2, 2**2, 3**2]})

In [4]: df.query('sqrt(a) >= 2')
---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
File ~/.local/lib/python3.9/site-packages/cudf/core/dataframe.py:7559, in extract_col(df, col)
7558 try:
-> 7559 return df._data[col]
7560 except KeyError:

File ~/.local/lib/python3.9/site-packages/cudf/core/column_accessor.py:155, in ColumnAccessor.__getitem__(self, key)
154 def __getitem__(self, key: Any) -> ColumnBase:
--> 155 return self._data[key]

KeyError: 'sqrt'

During handling of the above exception, another exception occurred:

KeyError Traceback (most recent call last)
Cell In [4], line 1
----> 1 df.query('sqrt(a) >= 2')

File ~/.local/lib/python3.9/site-packages/cudf/core/dataframe.py:4174, in DataFrame.query(self, expr, local_dict)
4168 callenv = {
4169 "locals": callframe.f_locals,
4170 "globals": callframe.f_globals,
4171 "local_dict": local_dict,
4172 }
4173 # Run query
-> 4174 boolmask = queryutils.query_execute(self, expr, callenv)
4175 return self._apply_boolean_mask(boolmask)

File ~/.local/lib/python3.9/site-packages/cudf/utils/queryutils.py:214, in query_execute(df, expr, callenv)
211 columns = compiled["colnames"]
213 # prepare col args
--> 214 colarrays = [cudf.core.dataframe.extract_col(df, col) for col in columns]
216 # wait to check the types until we know which cols are used
217 if any(col.dtype not in SUPPORTED_QUERY_TYPES for col in colarrays):

File ~/.local/lib/python3.9/site-packages/cudf/utils/queryutils.py:214, in (.0)
211 columns = compiled["colnames"]
213 # prepare col args
--> 214 colarrays = [cudf.core.dataframe.extract_col(df, col) for col in columns]
216 # wait to check the types until we know which cols are used
217 if any(col.dtype not in SUPPORTED_QUERY_TYPES for col in colarrays):

File ~/.local/lib/python3.9/site-packages/cudf/core/dataframe.py:7567, in extract_col(df, col)
7561 if (
7562 col == "index"
7563 and col not in df.index._data
7564 and not isinstance(df.index, MultiIndex)
7565 ):
7566 return df.index._data.columns[0]
-> 7567 return df.index._data[col]

File ~/.local/lib/python3.9/site-packages/cudf/core/column_accessor.py:155, in ColumnAccessor.__getitem__(self, key)
154 def __getitem__(self, key: Any) -> ColumnBase:
--> 155 return self._data[key]

KeyError: 'sqrt'

In [5]: df.to_pandas().query('sqrt(a) >= 2')
Out[5]:
a
1 4
2 9
```

**Additional context**
https://github.com/pandas-dev/pandas/blob/v1.5.2/pandas/core/computation/ops.py#L39-L60

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.