Skip to main content

A realistic workload: parallel pandas

Everything so far, applied to something recognizable: a dataframe pipeline that takes too long.

The starting point​

import pandas as pd

df = pd.read_parquet("s3://bucket/events/")     # 400M rows
df["bucket"] = df["value"].apply(classify)      # slow, row-wise
summary = df.groupby(["day", "bucket"]).agg(
    events=("id", "count"),
    total=("value", "sum"),
)

Forty minutes on one machine, most of it in apply.

The pattern​

Partition → distribute → combine. Split the input into independently processable pieces, run the per-piece work as tasks, combine the results.

Partition​

The data is already partitioned as files. Use them:

import eugo.hpc
from mylib.io import list_partitions          # returns object-storage paths

paths = list_partitions("s3://bucket/events/")
print(len(paths))                              # 512

512 partitions across 32 workers gives the scheduler 16 per worker, enough to keep everyone busy and absorb uneven partition sizes.

Partition in the data, not in memory

Reading 400M rows into your session to slice it up defeats the purpose: the session becomes the bottleneck and its memory the ceiling. Have each worker read its own partition.

Distribute​

@eugo.hpc.distribute
def process_partition(path):
    df = pd.read_parquet(path)                       # worker reads it
    df["bucket"] = df["value"].apply(classify)
    return df.groupby(["day", "bucket"]).agg(        # partial aggregate
        events=("id", "count"),
        total=("value", "sum"),
    )

The body is unchanged pandas. Nothing inside knows it runs on a cluster.

It returns a partial aggregate, not the raw rows. A grouped summary is orders of magnitude smaller than the partition it came from, so very little crosses the network.

Combine​

futures = [process_partition(p) for p in paths]
partials = eugo.hpc.get(futures)

summary = (
    pd.concat(partials)
    .groupby(["day", "bucket"])
    .agg(events=("events", "sum"), total=("total", "sum"))
)

Note the second aggregation. Concatenating partial groupbys leaves duplicate keys, the same (day, bucket) appearing once per partition, so they must be re-aggregated.

Which aggregations combine cleanly​

This is the part that bites.

Safe: count, sum, min, max. Combining partials is the same operation again.

Needs care: mean. Averaging averages is wrong when partitions differ in size. Carry sum and count separately, then divide at the end:

# In the task
.agg(total=("value", "sum"), n=("value", "count"))

# After combining
summary["mean"] = summary["total"] / summary["n"]

Not combinable: median, nunique, exact quantiles. These need the whole distribution. Either compute them on a sample, use an approximate algorithm, or repartition by the grouping key so each key lives entirely within one partition.

What to expect​

This is a mixed compute-and-I/O workload: reads from object storage, real computation in apply, small results back.

At 32 workers, expect well short of 32x — reading is a shared resource, and the pd.concat step is serial, so a long job comes down by a large but sub-linear factor.

The speedup explorer shows the shape of that falloff. Its figures are illustrative rather than measured, so use them to reason about the regime and measure your own workload for a number you can plan against.

When this pattern does not apply​

Cross-partition dependencies. A rolling window spanning partition boundaries needs overlapping reads, or repartitioning by the ordering key first.

Sorting the whole dataset. Not naturally partitionable. Sort within partitions and merge, or repartition by range.

Joins against another large table. Every partition needs matching rows from the other side. If one side is small, pass it as an argument or hold it in an actor; if both are large, this needs a shuffle and is a harder problem than fan-out.

The checklist​

  1. Partitions readable independently, by workers, from object storage
  2. Task bodies measured in seconds
  3. Small returns: aggregates or paths, not raw frames
  4. Combine step aware of which aggregations are safe
  5. A baseline timing to compare against

Video: Parallel pandas in practice. The same material, with a transcript.