Discussion about pandas.groupby, GroupBy.apply, GroupBy.agg and GroupBy.transform

All about .groupby

Let’s start with a pandas.DataFrame with some random data

  • time is a pandas.date_range from July 4, 2021 9:00 AM ~ 1:00 PM, with some random number of minutes add to put noise on each measurement
  • ticker is a randomly selected four-character string
  • price is a random starting value on [0, 100) plus some random noise drawn uniformly from [-5, +5)
from pandas import DataFrame, date_range, to_timedelta, concat
from numpy import repeat, array, tile
from numpy.random import choice, random, randint
from random import randrange
from string import ascii_lowercase
from datetime import timedelta

df = DataFrame({
  'time': repeat(date_range('2021-07-04 9:00', periods=(size:=5), freq='1H'), (rpt:=4)) + to_timedelta(randint(0, 59, size=(size*rpt)), unit='T'),
  'ticker': tile(choice([*ascii_lowercase], size=(size, 4)).view('<U4').ravel(), rpt),
  'price': tile(random(size)*100, rpt) + ((random(size*rpt) - .5) * 10),
  'signal': randint(10, size=size*rpt)
})
print(df.head(3))

We can use the .groupby method on this to perform operations on groups of ticker. In this case, .groupby method returns a DataFrameGroupby object that we can call further methods on to operate on the groups.

_print = lambda expr: print(f'{expr} '.center(50, '-'), eval(expr), sep='\n')
# print(f'{df.groupby("ticker") = }')
_print('df.groupby("ticker").sum().head(2)')

The pandas.DataFrameGroupBy object defines a set of operations that can be performed across on all of the groups, some of which are implemented as high-performance vectorized operations.

Some simple operations supported are:

  • .count
  • .min
  • .max
  • .sum
  • .var
  • .std
  • .size

These perform the specified reduction/aggregation operation on all relevant columns. If there are non-numeric columns, some operations may NOT be performed: e.g. .sum() will sum within the groups for all numeric columns and skip non=numeric columns; .size will return a single column with the size of the group; .count will count non-NaN values for all columns; .min and .max will find the minimum and maximum value for all comparable columns.

The return value will be a new pandas.Series or pandas.DataFrame where the data is the result of the desired operation and the index is the group values (potentially MultiIndex when grouping multiple values). A pandas.DataFrame if there are multiple columns resulting from the desired operation; A pandas.Series if there is only column.

for meth in 'count min max sum var std size'.split():
  _print(f'df.groupby("ticker").{meth}().head(2)')

With the .as_index=False flag, you can suppress the assignment of the group labels to the index.

_print = lambda expr: print('-'*50, f'{expr} '.center(50, '-'), '-'*50, eval(expr), sep='\n')

from pandas import DataFrame
from numpy.random import randint, choice, normal
from string import ascii_lowercase, ascii_uppercase

index = choice([*ascii_lowercase], size=(size:=10, 2)).view('<U2').ravel()
df = DataFrame({
  'x': randint(0, 10, size=size),
  'y': randint(0, 10, size=size),
}, index=index).sort_index()

_print('df.sample(2)')
_print("df.groupby(level=0).sum().head()")
_print("df.groupby(level=0, as_index=False).sum().head()")

Note that the .groupby has an axis= parameter, while we typically group down rows, it is also possible to group across columns (this is might be useful for joining similarly named columns). For the rest of the notes, we’ll ignore this possibility and assume we only wang grouping of rows (i.e. axis=0).

_print = lambda expr: print(f'{expr} '.center(50, '-'), eval(expr), sep='\n')

df = DataFrame({
  'time': repeat(date_range('2021-07-04 9:00', periods=(size:=5), freq='1H'), (rpt:=4)) + to_timedelta(randint(0, 59, size=(size*rpt)), unit='T'),
  'ticker': tile(choice([*ascii_lowercase], size=(size, 4)).view('<U4').ravel(), rpt),
  'price': tile(random(size)*100, rpt) + ((random(size*rpt) - .5) * 10),
  'sig1': randint(10, size=size*rpt),
  'sig2': randint(10, size=size*rpt),
  'sig3': randint(10, size=size*rpt),
})
df.columns = *df.columns[:-3], 'sig', 'sig', 'sig'

_print('df.groupby(level=0, axis=1).sum().head(3)')

We can perform a .groupby grouping by a single column, by multiple columns, or on the index. A considerable amount of time for a .groupby operation is spent sorting the data. If the data is pre-sorted, then we can disable this sorting with a sort=False flag at significant performance gain.

There is also a pandas.Series.GroupBy object in addition to the pabdas.DataFrameGroupBy. This represents a .groupby operation being performed on a single pandas.Series of data

index = choice([*ascii_lowercase], size=(size:=10, 2)).view('<U2').ravel()
df = DataFrame({
  'x': normal(size=size),
  'y': normal(size=size),
}, index=index).sort_index()

print(f'{df["x"].groupby(level=0) = }')

.groupby can also group by a piece of data that is NOT contained in the DataFrame or Series.

index = choice([*ascii_lowercase], size=(size:=10, 4)).view('<U4').ravel()
df = DataFrame({
  'x': randint(0, 10, size=size),
  'y': randint(0, 10, size=size),
}, index=index).sort_index()

_print("df.groupby(df['x'] < df['y']).sum()")
_print("df.groupby(lambda idx_val: idx_val.lower()[:2]).sum().head(2)")
_print("df.groupby(df.index.str.lower().str[:2]).sum().head(2)")

If we were to .groupby two fields, and we want the result to be non-“flat”; in other words, we want one of the fields to be the row index, the other field to be the column index, and the grouped function values to be at the intersection of these, then we can do a .groupby().unstack().

There is also a short-named for this operation: .pivot_table() which has the advantage of NOT converting the data values to a floating point if there are missing intersections fo the two groups.

_print = lambda expr: print('-'*50, f'{expr} '.center(50, '-'), '-'*50, eval(expr), sep='\n')

from pandas import DataFrame, MultiIndex
from sys import maxsize

index = MultiIndex.from_tuples([
  ('a', 'aa', 'aaa'),
  ('a', 'bb', 'aaa'),
  ('a', 'bb', 'abb'),
  ('b', 'aa', 'aaa'),
  ('b', 'bb', 'aaa'),
  ('b', 'cc', 'aaa'),
], names='A B C'.split())
df = DataFrame({
  'x': randint(0, 10, size=index.size),
  'y': 2**53 +1 # precision will be lost when converting to float type
}, index=index)

_print('df.groupby(level=(1, 2)).count().unstack()')
_print('df.groupby(level=(1, 2)).count().unstack().fillna(0).convert_dtypes()')
_print("df.pivot_table(index='B', columns='C', values=['x', 'y'], aggfunc='count', fill_value=0)")

All about .groupby().{aggregate, apply, filter, transform}

Going back to our data from before, we may want to perform a more complex operation on the gourds than a simple .count or .sum.

We can iterate over the groups directly and perform operations in Python (though this should be reserved for cases where the number of groups is on the order of the program structure, NOT on the order of the program computation). We can also use the .apply, .aggregate, .filter, and.transform methods. These largely fit into the following categories:

  1. aggregation (reducing the size of the group)
  2. transformation (preserving the indexing of the group)
  3. filtration (reducing the size of the group)

Let’s start with filtration, using .filter:

df = DataFrame({
  'time': repeat(date_range('2021-07-04 9:00', periods=(size:=5), freq='1H'), (rpt:=4)) + to_timedelta(randint(0, 59, size=(size*rpt)), unit='T'),
  'ticker': tile(choice([*ascii_lowercase], size=(size, 4)).view('<U4').ravel(), rpt),
  'price': tile(random(size)*100, rpt) + ((random(size*rpt) - .5) * 10),
  'signal': randint(10, size=size*rpt)
}).set_index(['ticker', 'time']).sort_index()

_print("df.groupby('ticker').filter(lambda g: (g['signal'].max() - g['signal'].min()) < 6).head(10)")

In a transformation, the result must be “like-indexed”. Let’s take a look at using .transform. The function using for .transform is restricted in the following ways:

  • It must return a value with the same shape as the input group or something that can be broadcast to the same shape
  • It must support being applied on a column-by-column basis
    • If it supports being applied to the entire group, then a fast path can be triggered
  • It must not change the groups in-place
_print("df.groupby('ticker').transform(lambda g: g.round(-1)).head(10)")
_print("df.groupby('ticker').transform(lambda g: g.round(-1).iloc[0]).head(10)")
_print("df.groupby('ticker').transform(lambda g: g.round(-1).iloc[:-1]).head(10)") # this will fail because of shape mismatch
_print("df.groupby('ticker').transform(lambda g: round(g, -1 if g.name == 'signal' else -2)).head(10)")

If we want to perform an aggregation, we can use .aggregate or .agg:

_print("df.groupby('ticker').aggregate('sum').head(10)")
_print("df.groupby('ticker').aggregate({'price': 'max', 'signal': ['min', 'max'], }).head(10)")
_print("df.groupby('ticker').aggregate(max_price = NamedAgg(column='price', aggfunc='max'), min_signal = ('signal', 'min')).head(10)")
_print("df.groupby('ticker').aggregate(lambda g: g.sum()).head(10)")

Finally, we have another means by which we can do aggregation operations: .apply. .apply is extremely flexible, but it has the downside of much slower in practice than .aggregate or .transform. .apply takes a function which accepts a DataFrame as its argument, returning a new DataFrame; the .apply machinery determines how to combine the result DataFrames into a a new structure.

_print("df.groupby('ticker').apply(lambda df: concat([df, df]).sort_index()).head(10)")
_print("df.groupby('ticker').apply(lambda df: (df['signal'] + df['price']).sum()).head(10)")

What about .rolling

The .groupby method is NOT the only method which can operate on the pandas.DataFrame._data region. The .rolling and .expanding operations provide similar functionality, but on rolling windows and for expanding transformations:

_print("df.rolling(10, min_periods=1).sum().head(10)")
_print("df.groupby('ticker').transform(lambda df: df.rolling(10, min_periods=1).sum()).head(10)")
_print("df.groupby('ticker')['price'].transform(lambda s: s.reset_index(level=0, drop=True).rolling(10).sum()).head(10)")

_print("df.expanding(min_periods=1).sum().head(10)")
_print("df.expanding(min_periods=1).count().head(10)")
_print("df.expanding(min_periods=1).aggregate(lambda s: s[-1] - s[0]).head(10)")
_print("df.expanding(min_periods=1).aggregate(lambda s: s.max() - s.min()).head(10)")

_print("df.ewm(alpha=0.1).mean().head(10)")
_print("df.ewm(alpha=0.1).std().head(10)")
_print("df.ewm(alpha=0.1).var().head(10)")