NVIDIA / NVIDIA/cudf

[BUG] JIT Groupby Apply `idxmax`/`idxmin` reductions return incorrect values when the data is all NaN

Open
#13,832 1 comment 0 reactions 0 assignees View on GitHub
bug numba Python
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Describe the bug**
When performing an `idxmax` or `idxmin` reduction on a dataframe column during JIT groupby apply, pandas returns `NaN` as the index label corresponding to the answer where as we return the index of the start of the group.

**Steps/Code to reproduce bug**
```python
import pandas as pd
import cudf
df = pd.DataFrame({
'a': [1, 1, 1, 2, 2, 2],
'b': [float('nan')] * 6
})

gdf = cudf.from_pandas(df, nan_as_null=False)

expect = df.groupby('a').apply(lambda x: x['b'].idxmax())
got = gdf.groupby('a').apply(lambda x: x['b'].idxmax(), engine='jit')

print(expect)
print(got)

```

```
a
1 NaN
2 NaN
dtype: float64
a
1 0
2 3
dtype: int64
```

**Expected behavior**
Ideally we'd match pandas.

**Environment overview (please complete the following information)**
Bare Metal, 23.10

**Additional context**
Originally came up [here](https://github.com/rapidsai/cudf/pull/11452#issuecomment-1403031678), and then again [here](https://github.com/rapidsai/cudf/pull/13820/files#r1284758449). This problem stems from the dtype of the answer being data dependent in pandas. In most cases, the `idx_{max,min}` functions return an `int64` if the index is of type `int64`, however this edge case of all `NaN` returns a `Nan` which is of float type. This poses a compatibility problem for the JIT engine as numba decides the types of all the variables in the input code up front, and currently an `idx_{max,min}` operation returns an `int64`. This leads to three options in my mind:

1. Return some kind of "sensible" int (current state). This leads to edge cases where our results differ from pandas.
2. Type `idxmax` and `idxmin` operations to return a float, e.g. cast the resulting integer to a float and return nan in the edge case, correctly. This trades a value mismatch for a dtype mismatch.
3. Raise in the edge case. This would require some engineering, related to https://github.com/rapidsai/cudf/issues/8774.

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.