Repository: dask/dask

Issue:
GroupBy get_group KeyError when category and npartitions >1
Appears there is an issue when using groupby get_group and the group was derived from a pandas categorical type. Only get error when npartitions is >1. Possibly related to #[5847](https://github.com/dask/dask/issues/5847)

````python
def checker(npartitions, toCategory):
    import dask.dataframe as dd
    import pandas as pd

    df = pd.DataFrame({'make':['a','b','c'], 'model': [2,5,7]})
    
    if toCategory:
        df['make'] = df['make'].astype('category')
        
    dd = dd.from_pandas(df, npartitions=npartitions)
    
    gbCol = 'make'
    for g in dd[gbCol].unique().compute():
        dd.groupby(gbCol).get_group(g).min(numeric_only=True).compute()
    return

checker(npartitions=2, toCategory=True)  # doesnt work, KeyError: 'a', 2 partitions
checker(npartitions=2, toCategory=False) # works
checker(npartitions=1, toCategory=True)  # works, only 1 partition
checker(npartitions=1, toCategory=False) # works
````
**Traceback**:
````python-traceback
Traceback (most recent call last):

  File "C:\Users\User1\AppData\Local\Temp\ipykernel_5608\2821702366.py", line 18, in <cell line: 18>
    checker(npartitions=2, toCategory=True)  # doesnt work, KeyError: 'a', 2 partitions

  File "C:\Users\User1\AppData\Local\Temp\ipykernel_5608\2821702366.py", line 15, in checker
    dd.groupby(gbCol).get_group(g).min(numeric_only=True).compute() #get their associated group which is dataframe type

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\base.py", line 315, in compute
    (result,) = compute(self, traverse=False, **kwargs)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\base.py", line 598, in compute
    results = schedule(dsk, keys, **kwargs)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\threaded.py", line 89, in get
    results = get_async(

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\local.py", line 511, in get_async
    raise_exception(exc, tb)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\local.py", line 319, in reraise
    raise exc

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\local.py", line 224, in execute_task
    result = _execute_task(task, data)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\optimization.py", line 990, in __call__
    return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 149, in get
    result = _execute_task(task, cache)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 119, in <genexpr>
    return func(*(_execute_task(a, cache) for a in args))

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 113, in _execute_task
    return [_execute_task(a, cache) for a in arg]

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 113, in <listcomp>
    return [_execute_task(a, cache) for a in arg]

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 119, in <genexpr>
    return func(*(_execute_task(a, cache) for a in args))

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 113, in _execute_task
    return [_execute_task(a, cache) for a in arg]

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 113, in <listcomp>
    return [_execute_task(a, cache) for a in arg]

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\utils.py", line 71, in apply
    return func(*args, **kwargs)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\dataframe\core.py", line 6650, in apply_and_enforce
    df = func(*args, **kwargs)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\dask\dataframe\groupby.py", line 239, in _groupby_get_group
    return grouped.get_group(get_key)

  File "C:\ProgramData\Anaconda3\envs\env3\lib\site-packages\pandas\core\groupby\groupby.py", line 747, in get_group
    raise KeyError(name)

KeyError: 'a'
````

**Environment**:

- Dask version: 2022.8.0
- Python version: 3.10
- Operating System: Windows 10
- Install method (conda, pip, source): conda


Discussion:
@ba05 Thanks for reporting! I can reproduce this, and it does look like a bug.

Note that the exact code above will NOT fail in previous versions of Dask because the `npartitions` calculated will be 1 (instead of 2) prior to `2022.8.0`. (the `npartitions` & `divisions` calculation changed in `2022.8.0`, ref: #9221)

We can use a bigger DataFrame to reproduce this error in previous Dask versions:

```python
import dask.dataframe as dd
import pandas as pd

df = pd.DataFrame(
    {"make": ["a", "b", "c", "d", "e", "f"], "model": [2, 5, 7, 9, 11, 13]}
)

df["make"] = df["make"].astype("category")

ddf = dd.from_pandas(df, npartitions=2)

ddf.groupby("make").get_group("a").compute() # KeyError: 'a'
```
The core issue is in how the categories are represented in Dask. The `make` column always retains metadata about *all* the categories, even if we wish to work only on a single partition.

Continuing from my above code snippet:

```python
for p in ddf.partitions:
    print(p.make.compute())

# Output:
# 
# 0    a
# 1    b
# 2    c
# Name: make, dtype: category
# Categories (6, object): ['a', 'b', 'c', 'd', 'e', 'f']
# 3    d
# 4    e
# 5    f
# Name: make, dtype: category
# Categories (6, object): ['a', 'b', 'c', 'd', 'e', 'f']
```

Since the metadata has all the groups, we see that the `else` block on line 242 is never reached:

https://github.com/dask/dask/blob/bf66221722cce8f09a9b09895bdb4596f14a5430/dask/dataframe/groupby.py#L233-L248

But, instead, the control always enters the `if` block and raises a `KeyError` when it can't find the `a` in the second partition which actually doesn't have `a`.

I'm not sure if having all the categories is a bug. If it is, we may need to look into `categorize()`.

If it isn't a bug, I think we can fix the original issue by updating the `if` check to something like:

```python
if df[by_key].isin([get_key]).any():
```

Resolving patch:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -234,12 +234,11 @@ def _groupby_get_group(df, by_key, get_key, columns):
     # SeriesGroupBy may pass df which includes group key
     grouped = _groupby_raise_unaligned(df, by=by_key)
 
-    if get_key in grouped.groups:
+    try:
         if is_dataframe_like(df):
             grouped = grouped[columns]
         return grouped.get_group(get_key)
-
-    else:
+    except KeyError:
         # to create empty DataFrame/Series, which has the same
         # dtype as the original
         if is_dataframe_like(df):
@@ -1633,7 +1632,10 @@ def last(self, split_every=None, split_out=1):
             token="last", func=M.last, split_every=split_every, split_out=split_out
         )
 
-    @derived_from(pd.core.groupby.GroupBy)
+    @derived_from(
+        pd.core.groupby.GroupBy,
+        inconsistencies="If the group is not present, Dask will return an empty Series/DataFrame.",
+    )
     def get_group(self, key):
         token = self._token_prefix + "get_group"
 


Test patch:
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -398,7 +398,8 @@ def test_groupby_multilevel_agg():
     assert_eq(res, sol)
 
 
-def test_groupby_get_group():
+@pytest.mark.parametrize("categoricals", [True, False])
+def test_groupby_get_group(categoricals):
     dsk = {
         ("x", 0): pd.DataFrame({"a": [1, 2, 6], "b": [4, 2, 7]}, index=[0, 1, 3]),
         ("x", 1): pd.DataFrame({"a": [4, 2, 6], "b": [3, 3, 1]}, index=[5, 6, 8]),
@@ -408,7 +409,15 @@ def test_groupby_get_group():
     d = dd.DataFrame(dsk, "x", meta, [0, 4, 9, 9])
     full = d.compute()
 
-    for ddkey, pdkey in [("b", "b"), (d.b, full.b), (d.b + 1, full.b + 1)]:
+    by_keys = [("b", "b"), (d.b, full.b)]
+
+    if categoricals:
+        d = d.categorize(columns=["b"])
+        full = d.compute()
+    else:
+        by_keys.append((d.b + 1, full.b + 1))
+
+    for ddkey, pdkey in by_keys:
         ddgrouped = d.groupby(ddkey)
         pdgrouped = full.groupby(pdkey)
         # DataFrame