pandas agg with custom aggregation function and dtype condition: column disappears silently

Question

I want to group and aggregate a pandas dataframe. I have a custom aggregation function that is supposed to do different things according to the dtype of the series it is applied to. However, some of the columns disappear silently. It seems to be linked to the dtype condition.

In the process of making a reproductible example, I noticed something weird:

if I test whether the series is of dtype string first, and then use an else statement for the other cases, my column with numbers is silently deleted
however, if I test series.dtype == "int64" first, it works as expected.

Here is my reproductible example:

import pandas as pd

# reprex dataframe
groups = ["group1", "group1", "group2", "group2"]
numbers = [1, 2, 3, 4]
letters = ["a", "b", "c", "c"]
data = {"group": groups, "number": numbers, "letter": letters}
df = pd.DataFrame(data)
df["letter"] = df["letter"].astype("string")

# custom aggregation functions
def is_int(series):
    if series.dtype == "int64":
        return "int64"
    else:
        return "not an int"

def is_string(series):
    if series.dtype == "string":
        return "string"
    else:
        return "not a string"

# the confusing part
df.groupby("group").agg(is_int) # works as expected
df.groupby("group").agg(is_string) # the "number" column disappears silently

Any insights on why this happens and how to deal with it are welcome!

user_4333359 · Answer

The new Dtypes are experimental, and as such there are alway weird quirks. The issue is that while a comparison for string dtypes with 'string' makes sense it fails for any other dtype.

df['letter'].dtype == 'string'
# True

df['number'].dtype == 'string'
#TypeError: data type 'string' not understood

So the groupby is silently ignoring the error in your second case and doesn't show an output because it's not able to aggregate that column. The correct comparison is with the actual string Dtype: pd.StringDtype, and use isinstance

def is_string(series):
    if isinstance(series.dtype, pd.StringDtype):
        return "string"
    else:
        return "not a string"

df.groupby("group").agg(is_string)
#              number  letter
#group                       
#group1  not a string  string
#group2  not a string  string

pandas agg with custom aggregation function and dtype condition: column disappears silently

1 Answers