Pandas Performance Optimization

Real-World Scenario

A data pipeline processing 10M rows of event logs takes 45 minutes with iterrows()-based logic. After vectorization, efficient dtypes, and chunked I/O, it runs in 3 minutes. The same pipeline on Polars runs in 25 seconds.

Dtype Optimization

import pandas as pd
import numpy as np

# Generate a realistic dataset
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),
    "category":  rng.choice(["electronics", "clothing", "food", "sports", "home"], n),
    "amount":    rng.uniform(0, 500, n).round(2),
    "quantity":  rng.integers(1, 10, n),
    "rating":    rng.integers(1, 6, n).astype(np.int8),  # 1-5 fits in int8
    "is_mobile": rng.integers(0, 2, n).astype(bool),
})

print(f"Before optimization: {df.memory_usage(deep=True).sum() / 1024**2:.1f} MB")

def optimize_dtypes(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()

    for col in df.select_dtypes("int64").columns:
        col_min, col_max = df[col].min(), df[col].max()
        for dtype in [np.int8, np.int16, np.int32]:
            if np.iinfo(dtype).min <= col_min and col_max <= np.iinfo(dtype).max:
                df[col] = df[col].astype(dtype)
                break

    for col in df.select_dtypes("float64").columns:
        if df[col].between(-3.4e38, 3.4e38).all():
            df[col] = df[col].astype(np.float32)

    for col in df.select_dtypes("object").columns:
        n_unique = df[col].nunique()
        if n_unique / len(df) < 0.5:   # < 50% unique values → use category
            df[col] = df[col].astype("category")

    return df

df_opt = optimize_dtypes(df)
print(f"After optimization:  {df_opt.memory_usage(deep=True).sum() / 1024**2:.1f} MB")

# Show dtypes after
print("\nDtype changes:")
for col in df.columns:
    if df[col].dtype != df_opt[col].dtype:
        print(f"  {col}: {df[col].dtype} → {df_opt[col].dtype}")

Vectorization vs iterrows

import pandas as pd
import numpy as np
import timeit

rng = np.random.default_rng(42)
df  = pd.DataFrame({
    "price":    rng.uniform(10, 500, 100_000),
    "quantity": rng.integers(1, 20, 100_000),
    "discount": rng.uniform(0, 0.3, 100_000),
    "category": rng.choice(["A", "B", "C"], 100_000),
})

# ── SLOW: iterrows ────────────────────────────────────────────────────
def slow_revenue(df: pd.DataFrame) -> pd.Series:
    revenues = []
    for _, row in df.iterrows():
        rev = row["price"] * row["quantity"] * (1 - row["discount"])
        if row["category"] == "A":
            rev *= 1.1
        revenues.append(rev)
    return pd.Series(revenues)

# ── FAST: vectorized ──────────────────────────────────────────────────
def fast_revenue(df: pd.DataFrame) -> pd.Series:
    base = df["price"] * df["quantity"] * (1 - df["discount"])
    return base * np.where(df["category"] == "A", 1.1, 1.0)

# Benchmark
t_slow = timeit.timeit(lambda: slow_revenue(df.head(10_000)), number=1)
t_fast = timeit.timeit(lambda: fast_revenue(df), number=10) / 10

print(f"iterrows (10k rows):    {t_slow:.3f}s")
print(f"vectorized (100k rows): {t_fast:.4f}s")
print(f"Speedup: >{t_slow / t_fast * 10:.0f}x")

# Verify same results
result_slow = slow_revenue(df.head(100))
result_fast = fast_revenue(df.head(100))
assert np.allclose(result_slow.values, result_fast.values, rtol=1e-5)
print("Results match ✓")

Efficient String Operations

import pandas as pd
import numpy as np
import re

rng = np.random.default_rng(42)
n   = 500_000

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

# ── SLOW: Python apply ────────────────────────────────────────────────
import timeit
def slow_domain(s: pd.Series) -> pd.Series:
    return s.apply(lambda x: x.split("@")[1].split(".")[0])

# ── FAST: .str accessor (vectorized C operations) ─────────────────────
def fast_domain(s: pd.Series) -> pd.Series:
    return s.str.split("@").str[1].str.split(".").str[0]

# Even faster: single regex extract
def fastest_domain(s: pd.Series) -> pd.Series:
    return s.str.extract(r"@(\w+)\.", expand=False)

t_slow    = timeit.timeit(lambda: slow_domain(emails.head(10_000)), number=1)
t_fast    = timeit.timeit(lambda: fast_domain(emails), number=5) / 5
t_fastest = timeit.timeit(lambda: fastest_domain(emails), number=5) / 5

print(f"apply (10k):       {t_slow:.3f}s")
print(f".str chain (500k): {t_fast:.3f}s")
print(f"str.extract (500k): {t_fastest:.3f}s")

Chunked Processing for Large Files

import pandas as pd
import numpy as np
from pathlib import Path

def process_large_csv(
    input_path: str,
    output_path: str,
    chunk_size: int = 100_000,
) -> dict:
    """Process a large CSV in chunks — never loads the full file into memory."""
    stats = {"rows_processed": 0, "rows_kept": 0, "chunks": 0}
    first_chunk = True

    for chunk in pd.read_csv(
        input_path,
        chunksize=chunk_size,
        dtype={"user_id": np.int32, "amount": np.float32},  # specify dtypes for efficiency
        parse_dates=["timestamp"],
    ):
        # Process: filter + aggregate per chunk
        chunk = chunk[chunk["amount"] > 0]
        chunk["revenue"] = chunk["amount"] * chunk.get("quantity", 1)

        # Write chunk to output (append after first chunk)
        chunk.to_csv(
            output_path,
            mode="w" if first_chunk else "a",
            header=first_chunk,
            index=False,
        )

        stats["rows_processed"] += len(chunk) + (chunk["amount"] <= 0).sum()
        stats["rows_kept"]      += len(chunk)
        stats["chunks"]         += 1
        first_chunk = False

    return stats


# Demo with synthetic data
def create_sample_csv(path: str, n: int = 1_000_000):
    rng = np.random.default_rng(42)
    df  = pd.DataFrame({
        "user_id":   rng.integers(1, 10_000, n).astype(np.int32),
        "amount":    (rng.uniform(-10, 500, n)).astype(np.float32),
        "quantity":  rng.integers(1, 10, n),
        "timestamp": pd.date_range("2024-01-01", periods=n, freq="s"),
    })
    df.to_csv(path, index=False)
    print(f"Created {path}: {n:,} rows")

Path("/tmp/large.csv").parent.mkdir(exist_ok=True)
create_sample_csv("/tmp/large.csv", 500_000)
stats = process_large_csv("/tmp/large.csv", "/tmp/output.csv", chunk_size=50_000)
print(f"Stats: {stats}")

When to Use Polars

# pip install polars
import polars as pl
import pandas as pd
import numpy as np
import timeit

# Generate data
rng = np.random.default_rng(42)
n   = 2_000_000
data = {
    "user_id":  rng.integers(1, 100_000, n),
    "event":    np.where(rng.random(n) > 0.5, "purchase", "view"),
    "amount":   rng.uniform(0, 500, n).round(2),
    "category": rng.choice(["A", "B", "C", "D"], n),
}

df_pd = pd.DataFrame(data)
df_pl = pl.DataFrame(data)

# Task: filter purchases, group by user and category, compute mean amount
def pandas_query(df: pd.DataFrame) -> pd.DataFrame:
    return (
        df[df["event"] == "purchase"]
        .groupby(["user_id", "category"])["amount"]
        .mean()
        .reset_index()
    )

def polars_query(df: pl.DataFrame) -> pl.DataFrame:
    return (
        df.filter(pl.col("event") == "purchase")
        .group_by(["user_id", "category"])
        .agg(pl.col("amount").mean())
    )

t_pd = timeit.timeit(lambda: pandas_query(df_pd), number=3) / 3
t_pl = timeit.timeit(lambda: polars_query(df_pl), number=3) / 3

print(f"Pandas:  {t_pd:.3f}s")
print(f"Polars:  {t_pl:.3f}s")
print(f"Speedup: {t_pd/t_pl:.1f}x")

# Polars lazy evaluation (query planning optimization)
result = (
    df_pl.lazy()
    .filter(pl.col("event") == "purchase")
    .group_by(["user_id", "category"])
    .agg([
        pl.col("amount").mean().alias("avg_amount"),
        pl.col("amount").sum().alias("total_amount"),
        pl.col("amount").count().alias("n_purchases"),
    ])
    .filter(pl.col("n_purchases") >= 3)   # only users with 3+ purchases
    .sort("total_amount", descending=True)
    .limit(10)
    .collect()   # execute the whole plan at once
)
print(f"\nTop 10 users by total spend:\n{result}")