GroupBy nth with Categorical, NA Data and dropna

[WillAyd]

The nth method of GroupBy objects can accept a dropna keyword, but it doesn't handle Categorical data appropriately. For instance:

In [9]: cat = pd.Categorical(['a', np.nan, np.nan], categories=['a', 'b'])
   ...: ser = pd.Series([1, 2, 3])
   ...: df = pd.DataFrame({'cat': cat, 'ser': ser})

In [10]: df.groupby('cat')['ser'].nth(0, dropna='all')
ValueError: Length mismatch: Expected axis has 3 elements, new values have 2 elements

In [11]: df.groupby('cat', observed=True)['ser'].nth(0, dropna='all')
ValueError: Length mismatch: Expected axis has 3 elements, new values have 1 elements

Note that you can get a nonsensical result if the length of the categories matches the length of the Categorical:

In [9]: cat = pd.Categorical(['a', np.nan, np.nan], categories=['a', 'b', 'c'])
   ...: ser = pd.Series([1, 2, 3])
   ...: df = pd.DataFrame({'cat': cat, 'ser': ser})
In [10]: df.groupby('cat')['ser'].nth(0, dropna='all')
Out[12]:
cat
a    1
b    2
c    3
Name: ser, dtype: int64

I'm not sure it makes sense to specify both observed and dropna anyway, so I think we should be raising if a categorical with NA values is detected within that method

@TomAugspurger

[jorisvandenbossche]

This doesn't seem specific to Categorical data, I get exactly the same error with object data:

In [27]: df['cat'] = df['cat'].astype(object)                                                                                                                                                                       

In [28]: df.dtypes                                                                                                                                                                                                  
Out[28]: 
cat    object
ser     int64
dtype: object

In [29]: df.groupby('cat')['ser'].nth(0, dropna='all')                                                                                                                                                              
...
ValueError: Length mismatch: Expected axis has 3 elements, new values have 1 elements

[WillAyd]

@jorisvandenbossche great catch! So it looks like handling NA values in the grouping is rather suspect, though worth calling out that is specific only to SeriesGroupBy:

>>> df.groupby('cat').nth(0, dropna='all')
     ser
cat
a      1

So might be two separate issues here

[WillAyd]

May or may not be intertwined with #24880