Dask Partition Parallel Aggregation

Scale pandas aggregations with Dask partitions, letting the scheduler parallelize map-side work and tree-reduce the shuffled results.

PythonResilience

Code

snippet
import dask.dataframe as dd
import pandas as pd

base = pd.DataFrame({
    'sensor': [f's{i % 8}' for i in range(100_000)],
    'reading': [round(20 + (i % 37) * 0.5, 2) for i in range(100_000)],
    'hour': [(i // 3600) % 24 for i in range(100_000)],
})

ddf = dd.from_pandas(base, npartitions=8)
print('partition count:', ddf.npartitions)


def peak_of(partition: pd.DataFrame) -> pd.DataFrame:
    idx = partition['reading'].idxmax()
    return partition.loc[[idx]]


daily = (
    ddf
    .groupby('hour')
    .agg(
        mean_reading=('reading', 'mean'),
        max_reading=('reading', 'max'),
        samples=('reading', 'count'),
    )
    .compute()
)

peaks = (
    ddf
    .map_partitions(peak_of, meta=base.head(0))
    .compute()
    .nlargest(3, 'reading')
)

print('hours reported:', len(daily))
print('global top readings:')
print(peaks[['sensor', 'reading']].to_string(index=False))