[Python] pandas group by on a subset returns wrong number of rows

I am dramatically simplifying this question (compared to the initial draft).

Given a dataframe with a categorical type, if I subset the dataframe so it only contains a few of the categorical values, the subsetted dataframe still remembers the full list of category levels.

When I groupby on this, smaller, dataframe, I continue to get values which I don't expect.

Here is some code which replicates this:

import pandas as pd

df = pd.DataFrame({
'colA':['a', 'b', 'b', 'c', 'c', 'c'],
'colB':1
})
df['colA'] = df['colA'].astype('category') #<= this is important
df


output

colA colB
0 a 1
1 b 1
2 b 1
3 c 1
4 c 1
5 c 1


df.groupby('colA').max()


output

colB
colA
a 1
b 1
c 1


smalldf = df.iloc[:3]
smalldf['colA']


output

0 a
1 b
2 b
Name: colA, dtype: category
Categories (3, object): ['a', 'b', 'c']


Notice above that 'c' is still in the category levels

smalldf.groupby('colA').max()


output

colB
colA
a 1.0
b 1.0
c NaN


This 'c' here is unexpected. However, I suppose it can be explained away because subsetting doesn't remove the category levels.

The problem gets worse when you write out the smaller dataframe

FEATHER_PATH = Path.joinpath(Path.home(), 'Downloads', 'test.feather')
smalldf.to_feather(FEATHER_PATH)

pd.read_feather(FEATHER_PATH)['colA']


0 a
1 b
2 b
Name: colA, dtype: category
Categories (3, object): ['a', 'b', 'c']


Say you write out the smaller dataframe and pass it to someone. They take a look at the file and only see colA values of 'a' and 'b'. However, when they do a groupby, 'c' shows up. 'c' is no where to be found in the actual data!

They would have to explicitly look for category levels to find the 'c' hidden in there!

I'm not sure if this behavior is wrong, but can certainly can be extremely confusing!

Not sure if this can be left here as a warning to others.

Continue reading...