What is the best way to retain data for columns that can't be aggregated during a a groupby operation on a pandas dataframe?

Viewed 45

Pandas dataframes are great for storing non-homogeneous data types, and the groupby method adds a lot of flexibility for aggregating the data; however, I often would like to keep underlying data that can't be aggregated after performing the groupby operation.

Example:

Let's work with the following dataframe:

df_cars = pd.DataFrame({
          'brand': ['Porsche', 'Porsche', 'Porsche', 'BMW' ,'BMW' ,'BMW'],
          'model': ['911 GT3', '911', '911', 'M3', 'M3', 'M3', ],
          'speed(mph)': [100, 120, 110, 70, 80, 90]})

If I were to run the following code:

df_cars.groupby('brand').mean()

I would get the mean speed(mph) on a per-brand basis; however, I would lose the underlying data from the model and brand columns.

In order to keep that information and perform the same mean calculation I would need to run the following code instead:

df_cars.groupby('brand').agg({
                        'speed(mph)':'mean', 
                        'model': lambda x: x.iloc[0] if x.nunique() == 1 else list(x),
                        'brand': lambda x: x.iloc[0] if x.nunique() == 1 else list(x)
                        })

This returns a dataframe with the mean calculated on a per-brand basis and keeps the underlying data (now collapsed) in the resulting dataframe as well.

Is this the most efficient way to retain this information? Furthermore, why is this not an option in the default implementation of groupby in order to more easily retain data that can't be aggregated?

1 Answers

First of all, when you aggregate data, in definition you will lose information, since youre summarizing your dataset.

That said, I see you have two options here:

1. simplify your code and apply correct methods:

df_cars.groupby('brand', as_index=False).agg({
    'speed(mph)': 'mean',
    'model': 'unique'
})

     brand  speed(mph)           model
0      BMW          80            [M3]
1  Porsche         110  [911 GT3, 911]

Or if you don't want your object type column in a list:

df_cars.groupby('brand', as_index=False).agg({
    'speed(mph)': 'mean',
    'model': lambda x: ','.join(np.unique(x))
})

     brand  speed(mph)        model
0      BMW          80           M3
1  Porsche         110  911,911 GT3

2. Split data by type of columns and apply simple aggregations:

this works for larger dataframes, with multiple object and numeric columns

df_cars = df_cars.set_index('brand')
text = df_cars.select_dtypes('object')
numeric = df_cars.select_dtypes(exclude='object')

g1 = text.groupby('brand')
g2 = numeric.groupby('brand')

dfn = g1.agg(np.unique).join(g2.mean()).reset_index()

     brand           model  speed(mph)
0      BMW              M3          80
1  Porsche  [911, 911 GT3]         110
Related