Pandas Performance Optimization

Real-World Scenario

A nightly job on event logs keeps getting slower as the data grows. The usual culprits are columns stored in wider types than they need, row-by-row Python loops, and loading a whole file when a streaming pass would do. Each of those can be measured and fixed.

Measuring and Shrinking Memory

memory_usage(deep=True) counts the real bytes, including the strings themselves.

import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
n = 1_000_000
df = pd.DataFrame({
    "user_id": rng.integers(1, 100_000, n),
    "event":   rng.choice(["click", "view", "purchase", "scroll"], n),
    "rating":  rng.integers(1, 6, n),
    "amount":  rng.uniform(0, 500, n),
})

small = df.assign(
    user_id=pd.to_numeric(df["user_id"], downcast="integer"),
    event=df["event"].astype("category"),
    rating=pd.to_numeric(df["rating"], downcast="integer"),
    amount=pd.to_numeric(df["amount"], downcast="float"),
)
before, after = df.memory_usage(deep=True), small.memory_usage(deep=True)
print((pd.DataFrame({"before_MB": before, "after_MB": after}) / 1e6).round(1))
print(small.dtypes)

$ python memory.py
         before_MB  after_MB
Index          0.0       0.0
user_id        8.0       4.0
event         54.7       1.0
rating         8.0       1.0
amount         8.0       4.0
user_id       int32
event      category
rating         int8
amount      float32
dtype: object

The event column dominates: a million separate strings become a million small integer codes plus four stored labels. downcast picks the smallest type that holds every value — int32 for IDs up to 100,000, int8 for ratings 1-5 — and float32 halves amount at the cost of precision, which is fine for most measurements but loses digits in large running totals.

select_dtypes and Pandas 3 Strings

Older dtype-optimizing helpers find text columns with select_dtypes("object").

import warnings
import pandas as pd

df = pd.DataFrame({"event": ["click", "view"], "rating": [4, 5]})

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    print(list(df.select_dtypes("object").columns))
print([w.category.__name__ for w in caught])

print(list(df.select_dtypes(include="str").columns))

$ python select_str.py
['event']
['Pandas4Warning']
['event']

In Pandas 3 the event column’s dtype is str, not object. select_dtypes("object") still returns it for compatibility, but warns that this will stop in a future version — at which point such helpers would silently skip every text column. include="str" says what you mean.

iterrows Versus itertuples Versus Vectorized

The same revenue calculation, three ways, timed on 10,000 rows.

import timeit
import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
df = pd.DataFrame({"price": rng.uniform(10, 500, 10_000),
                   "qty": rng.integers(1, 20, 10_000),
                   "cat": rng.choice(["A", "B"], 10_000)})

def with_iterrows():
    return [r["price"] * r["qty"] * (1.1 if r["cat"] == "A" else 1.0) for _, r in df.iterrows()]

def with_itertuples():
    return [r.price * r.qty * (1.1 if r.cat == "A" else 1.0) for r in df.itertuples()]

def vectorized():
    return df["price"] * df["qty"] * np.where(df["cat"] == "A", 1.1, 1.0)

print("same:", np.allclose(with_iterrows(), vectorized()))
for f in (with_iterrows, with_itertuples, vectorized):
    print(f"{f.__name__:16} {min(timeit.repeat(f, number=1, repeat=5)) * 1000:8.2f} ms")

$ python loops.py
same: True
with_iterrows      124.61 ms
with_itertuples      5.13 ms
vectorized           0.41 ms

All three agree on the answer; the difference is cost. Best of five runs on one machine: iterrows builds a full Series for every row, itertuples yields plain named tuples and is far cheaper, and the vectorized version does the whole column in a few array operations. Your timings will differ, but the ordering is the part to keep.

String Operations

With Pandas 3’s default string dtype, it is worth timing .str methods rather than assuming they beat apply.

import timeit
import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
emails = pd.Series([f"user{i}@{'gmail' if i % 3 == 0 else 'yahoo'}.com"
                    for i in rng.integers(1, 100_000, 200_000)])

ways = {
    "apply":       lambda: emails.apply(lambda x: x.split("@")[1].split(".")[0]),
    ".str chain":  lambda: emails.str.split("@").str[1].str.split(".").str[0],
    ".str.extract": lambda: emails.str.extract(r"@(\w+)\.", expand=False),
}
print(emails.dtype, emails.str.extract(r"@(\w+)\.", expand=False).value_counts().to_dict())
for name, f in ways.items():
    print(f"{name:13} {min(timeit.repeat(f, number=1, repeat=3)) * 1000:7.1f} ms")

$ python str_speed.py
str {'yahoo': 133775, 'gmail': 66225}
apply            39.7 ms
.str chain      138.2 ms
.str.extract     57.5 ms

This environment stores strings in the Python-backed str dtype (PyArrow is not installed), so each .str step still processes Python string objects. Compare the three lines above rather than assuming the vectorized-looking form wins: here the three-step .str chain was the slowest of the three, and even the single str.extract did not beat a plain apply. Installing PyArrow switches the default string storage to PyArrow, so re-time string-heavy code in whichever environment it will actually run in.

Processing a CSV in Chunks

chunksize turns read_csv into an iterator, so only one chunk is in memory at a time.

import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
n = 300_000
pd.DataFrame({"category": rng.choice(["A", "B", "C"], n),
              "amount": rng.integers(-10, 500, n)}).to_csv("events.csv", index=False)

totals, chunks = pd.Series(dtype="int64"), 0
for chunk in pd.read_csv("events.csv", chunksize=100_000):
    kept = chunk[chunk["amount"] > 0]
    totals = totals.add(kept.groupby("category")["amount"].sum(), fill_value=0)
    chunks += 1

full = pd.read_csv("events.csv").query("amount > 0").groupby("category")["amount"].sum()
print("chunks:", chunks)
print(totals.astype("int64"))
print("matches full read:", totals.astype("int64").equals(full))

$ python chunks.py
chunks: 3
category
A    24468012
B    24427689
C    24428927
dtype: int64
matches full read: True

Each chunk is filtered and summarized on its own, and the partial sums are combined with add(..., fill_value=0). The result is identical to reading the whole file at once, but peak memory is bounded by one 100,000-row chunk rather than the file size. This only works for aggregations that can be combined from pieces — sums and counts can, a median cannot.

Pandas Versus Polars on the Same Query

Polars is a separate DataFrame library (pip install polars) with a multi-threaded engine.

import timeit
import numpy as np
import pandas as pd
import polars as pl

rng = np.random.default_rng(42)
n = 2_000_000
data = {"user_id": rng.integers(1, 100_000, n),
        "event": rng.choice(["purchase", "view"], n),
        "amount": rng.uniform(0, 500, n)}
df_pd, df_pl = pd.DataFrame(data), pl.DataFrame(data)

def with_pandas():
    return df_pd[df_pd["event"] == "purchase"].groupby("user_id")["amount"].mean()

def with_polars():
    return df_pl.filter(pl.col("event") == "purchase").group_by("user_id").agg(pl.col("amount").mean())

check = with_polars().sort("user_id")["amount"].to_numpy()
print("same result:", np.allclose(with_pandas().to_numpy(), check))
t_pd = min(timeit.repeat(with_pandas, number=1, repeat=5))
t_pl = min(timeit.repeat(with_polars, number=1, repeat=5))
print(f"pandas {t_pd * 1000:.0f} ms   polars {t_pl * 1000:.0f} ms   ratio {t_pd / t_pl:.1f}x")

$ python polars_compare.py
same result: True
pandas 142 ms   polars 17 ms   ratio 8.6x

Both libraries produce the same per-user averages; Polars returns groups unsorted, so the check sorts first. The timings are best of five on one machine with Polars using all its cores, and the ratio depends heavily on the query and the hardware — measure your own bottleneck before rewriting a pipeline.

Practice

1. In memory.py, convert only the event column to category and print how many bytes that one column uses before and after.

import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
event = pd.Series(rng.choice(["click", "view", "purchase", "scroll"], 1_000_000))
cat = event.astype("category")
print(event.memory_usage(deep=True), cat.memory_usage(deep=True))
print(cat.cat.categories.tolist(), cat.cat.codes.dtype)

$ python practice_category.py
54751032 1000351
['click', 'purchase', 'scroll', 'view'] int8

2. In chunks.py, also count how many rows had a non-positive amount, across all chunks.

import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
pd.DataFrame({"amount": rng.integers(-10, 500, 300_000)}).to_csv("amounts.csv", index=False)

dropped = sum((chunk["amount"] <= 0).sum()
              for chunk in pd.read_csv("amounts.csv", chunksize=100_000))
print("non-positive rows:", dropped)
print("check:", (pd.read_csv("amounts.csv")["amount"] <= 0).sum())

$ python practice_chunks.py
non-positive rows: 6465
check: 6465

3. In loops.py, time df.apply(..., axis=1) for the same calculation and see where it lands relative to iterrows and itertuples.

import timeit
import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
df = pd.DataFrame({"price": rng.uniform(10, 500, 10_000),
                   "qty": rng.integers(1, 20, 10_000),
                   "cat": rng.choice(["A", "B"], 10_000)})

def rev(r):
    return r["price"] * r["qty"] * (1.1 if r["cat"] == "A" else 1.0)

t = min(timeit.repeat(lambda: df.apply(rev, axis=1), number=1, repeat=5))
print(f"apply(axis=1) {t * 1000:.2f} ms")

$ python practice_apply.py
apply(axis=1) 37.68 ms