Live Notebook

You can run this notebook in a live session Binder or view it on Github.

DataFrames: Groupby

This notebook uses the Pandas groupby-aggregate and groupby-apply on scalable Dask dataframes. It will discuss both common use and best practices.

Start Dask Client for Dashboard

Starting the Dask Client is optional. It will provide a dashboard which is useful to gain insight on the computation.

The link to the dashboard will become visible when you create the client below. We recommend having it open on one side of your screen while using your notebook on the other side. This can take some effort to arrange your windows, but seeing them both at the same is very useful when learning.

[1]:
from dask.distributed import Client
client = Client(n_workers=1, threads_per_worker=4, processes=False, memory_limit='2GB')
client
[1]:

Client

Cluster

  • Workers: 1
  • Cores: 4
  • Memory: 2.00 GB

Artifical dataset

We create an artificial timeseries dataset to help us work with groupby operations

[2]:
import dask
df = dask.datasets.timeseries()
df
[2]:
Dask DataFrame Structure:
id name x y
npartitions=30
2000-01-01 int64 object float64 float64
2000-01-02 ... ... ... ...
... ... ... ... ...
2000-01-30 ... ... ... ...
2000-01-31 ... ... ... ...
Dask Name: make-timeseries, 30 tasks

This dataset is small enough to fit in the cluster’s memory, so we persist it now.

You would skip this step if your dataset becomes too large to fit into memory.

[3]:
df = df.persist()

Groupby Aggregations

Dask dataframes implement a commonly used subset of the Pandas groupby API (see Pandas Groupby Documentation.

We start with groupby aggregations. These are generally fairly efficient, assuming that the number of groups is small (less than a million).

[4]:
df.groupby('name').x.mean().compute()
[4]:
name
Alice      -0.001095
Bob        -0.001230
Charlie    -0.000754
Dan         0.002367
Edith      -0.000065
Frank       0.000644
George     -0.000250
Hannah     -0.001574
Ingrid      0.001598
Jerry      -0.001824
Kevin       0.000242
Laura       0.000283
Michael    -0.001043
Norbert    -0.000622
Oliver      0.002783
Patricia    0.002519
Quinn       0.001060
Ray         0.003793
Sarah      -0.002240
Tim         0.003610
Ursula      0.000757
Victor     -0.003148
Wendy      -0.000516
Xavier     -0.001256
Yvonne     -0.000977
Zelda      -0.002392
Name: x, dtype: float64

Performance will depend on the aggregation you do (mean vs std), the key on which you group (name vs id), and the number of total groups

[5]:
%time _ = df.groupby('id').x.mean().compute()
CPU times: user 177 ms, sys: 12.9 ms, total: 190 ms
Wall time: 178 ms
[6]:
%time _ = df.groupby('name').x.mean().compute()
CPU times: user 486 ms, sys: 19.6 ms, total: 506 ms
Wall time: 478 ms
[7]:
%time df.groupby('name').agg({'x': ['mean', 'std'], 'y': ['mean', 'count']}).compute().head()
CPU times: user 391 ms, sys: 25.5 ms, total: 416 ms
Wall time: 397 ms
[7]:
x y
mean std mean count
name
Alice -0.001095 0.577039 -0.000867 99792
Bob -0.001230 0.577500 -0.001916 99381
Charlie -0.000754 0.577407 0.001138 99342
Dan 0.002367 0.576901 -0.001437 99199
Edith -0.000065 0.577902 0.000614 100000

This is the same as with Pandas. Generally speaking, Dask.dataframe groupby-aggregations are roughly same performance as Pandas groupby-aggregations, just more scalable.

You can read more about Pandas’ common aggregations in the Pandas documentation.

Custom Aggregations

Dask dataframe Aggregate is available for custom aggregations (See Dask dataframe Aggregate Documentation)

Many groups

By default groupby-aggregations (like groupby-mean or groupby-sum) return the result as a single-partition Dask dataframe. Their results are usually quite small, so this is usually a good choice.

However, sometimes people want to do groupby aggregations on many groups (millions or more). In these cases the full result may not fit into a single Pandas dataframe output, and you may need to split your output into multiple partitions. You can control this with the split_out= parameter

[8]:
# Computational graph of a single output aggregation (for a small number of groups, like 1000)
df.groupby('name').x.mean().visualize(node_attr={'penwidth': '6'})
Unable to revert mtime: /usr/share/fonts/cMap
Unable to revert mtime: /usr/share/fonts/cmap/adobe-cns1
Unable to revert mtime: /usr/share/fonts/cmap/adobe-gb1
Unable to revert mtime: /usr/share/fonts/cmap/adobe-japan1
Unable to revert mtime: /usr/share/fonts/cmap/adobe-japan2
Unable to revert mtime: /usr/share/fonts/cmap/adobe-korea1
[8]:
../_images/dataframes_02-groupby_16_1.png
[9]:
# Computational graph of an aggregation to four outputs (for a larger number of groups, like 1000000)
df.groupby('id').x.mean(split_out=4).visualize(node_attr={'penwidth': '6'})
Unable to revert mtime: /usr/share/fonts/cMap
Unable to revert mtime: /usr/share/fonts/cmap/adobe-cns1
Unable to revert mtime: /usr/share/fonts/cmap/adobe-gb1
Unable to revert mtime: /usr/share/fonts/cmap/adobe-japan1
Unable to revert mtime: /usr/share/fonts/cmap/adobe-japan2
Unable to revert mtime: /usr/share/fonts/cmap/adobe-korea1
[9]:
../_images/dataframes_02-groupby_17_1.png

Groupby Apply

Groupby-aggregations are generally quite fast because they can be broken down easily into well known operations. The data doesn’t have to move around too much and we can just pass around small intermediate values across the network.

For some operations however the function to be applied requires all data from a given group (like every record of someone named “Alice”). This will force a great deal of communication and be more expensive, but is still possible with the Groupby-apply method. This should be avoided if a groupby-aggregation works.

In the following example we train a simple Scikit-Learn machine learning model on every person’s name.

[10]:
from sklearn.linear_model import LinearRegression

def train(partition):
    est = LinearRegression()
    est.fit(partition[['x', 'id']].values, partition.y.values)
    return est
[11]:
%time df.groupby('name').apply(train, meta=object).compute().sort_index()
CPU times: user 3.6 s, sys: 438 ms, total: 4.04 s
Wall time: 3.87 s
[11]:
name
Alice       LinearRegression()
Bob         LinearRegression()
Charlie     LinearRegression()
Dan         LinearRegression()
Edith       LinearRegression()
Frank       LinearRegression()
George      LinearRegression()
Hannah      LinearRegression()
Ingrid      LinearRegression()
Jerry       LinearRegression()
Kevin       LinearRegression()
Laura       LinearRegression()
Michael     LinearRegression()
Norbert     LinearRegression()
Oliver      LinearRegression()
Patricia    LinearRegression()
Quinn       LinearRegression()
Ray         LinearRegression()
Sarah       LinearRegression()
Tim         LinearRegression()
Ursula      LinearRegression()
Victor      LinearRegression()
Wendy       LinearRegression()
Xavier      LinearRegression()
Yvonne      LinearRegression()
Zelda       LinearRegression()
dtype: object