Chapter 149 · Advanced & Applied Topics
Big Data & Scaling · Challenge Solutions
Worked solutions: shrink the memory footprint, aggregate a file in chunks, compare CSV with Parquet, write MapReduce by hand, and diagnose then repair data skew.
In [1]:
import numpy as np, pandas as pd, os, time, tempfile
import matplotlib.pyplot as plt
plt.rcParams.update({"figure.dpi":110,"axes.grid":True,"grid.alpha":0.25,"font.size":11})
FU, BL, GR, RD, AM = "#a21caf", "#2563eb", "#16a34a", "#dc2626", "#d97706"
BASE = "https://raw.githubusercontent.com/johnfisher-ai/Statistics-Data-Science-AI-Visual-Book/main/data/"
fn = "big-data-and-scaling--ride-logs.xlsx"
try:
df = pd.read_excel("../../data/" + fn)
except FileNotFoundError:
df = pd.read_excel(BASE + fn)
print(df.shape)
print(df.head(3).to_string())
# A scratch folder for the larger files we build below. Nothing here is written to the book repo.
TMP = tempfile.gettempdir()
print("\nscratch folder:", TMP)
(25000, 11) ride_id city vehicle_type pickup_hour distance_km duration_min surge_multiplier fare rider_id payment_type rating 0 1 Metro City standard 18 5.55 27.3 1.2 20.34 1499 card 5.0 1 2 Riverton standard 19 1.85 11.9 1.2 10.51 1575 card 4.7 2 3 Metro City standard 22 4.78 18.7 1.0 14.04 228 wallet 4.2 scratch folder: /var/folders/wv/rntn6xtd407cmdsyx0b0wwnw0000gn/T
Challenge 1 · Cut the memory footprint in half¶
Downcast the types and report the saving, then confirm the numbers did not move.
In [2]:
opt = df.copy()
for col in opt.select_dtypes("object"):
if opt[col].nunique() / len(opt) < 0.5: opt[col] = opt[col].astype("category")
for col in opt.select_dtypes("integer"): opt[col] = pd.to_numeric(opt[col], downcast="integer")
for col in opt.select_dtypes("float"): opt[col] = pd.to_numeric(opt[col], downcast="float")
b, a = df.memory_usage(deep=True).sum(), opt.memory_usage(deep=True).sum()
print(f"{b/1e6:.2f} MB -> {a/1e6:.2f} MB ({100*(1-a/b):.0f}% saved)")
print("mean fare unchanged:", round(float(df.fare.mean()), 4), "vs", round(float(opt.fare.mean()), 4))
2.70 MB -> 0.70 MB (74% saved) mean fare unchanged: 18.0374 vs 18.0374
/var/folders/wv/rntn6xtd407cmdsyx0b0wwnw0000gn/T/ipykernel_29720/923491595.py:2: Pandas4Warning: For backward compatibility, 'str' dtypes are included by select_dtypes when 'object' dtype is specified. This behavior is deprecated and will be removed in a future version. Explicitly pass 'str' to `include` to select them, or to `exclude` to remove them and silence this warning.
See https://pandas.pydata.org/docs/user_guide/migration-3-strings.html#string-migration-select-dtypes for details on how to write code that works with pandas 2 and 3.
for col in opt.select_dtypes("object"):
Challenge 2 · Average trip distance by vehicle type, in chunks¶
Never hold the whole file. Remember that means carrying sums and counts, not means.
In [3]:
big = pd.concat([df]*20, ignore_index=True)
p = os.path.join(TMP, "sol_rides.csv"); big.to_csv(p, index=False)
partials = [ch.groupby("vehicle_type").distance_km.agg(["sum", "count"]) for ch in pd.read_csv(p, chunksize=50_000)]
out = pd.concat(partials).groupby(level=0).sum()
out["avg_km"] = out["sum"] / out["count"]
print(out[["count", "avg_km"]].round(2).to_string())
print("\nmatches the full-file answer?",
np.allclose(out.avg_km.sort_index(), big.groupby("vehicle_type").distance_km.mean().sort_index()))
count avg_km vehicle_type eco 75460 4.24 premium 61460 4.28 standard 273580 4.33 xl 89500 4.33 matches the full-file answer? True
Challenge 3 · How much does Parquet actually buy you?¶
Compare size and the time to read a single column.
In [4]:
q = os.path.join(TMP, "sol_rides.parquet"); big.to_parquet(q, index=False)
cs, ps = os.path.getsize(p)/1e6, os.path.getsize(q)/1e6
t0 = time.perf_counter(); pd.read_csv(p, usecols=["fare"]); tc = time.perf_counter()-t0
t0 = time.perf_counter(); pd.read_parquet(q, columns=["fare"]); tp = time.perf_counter()-t0
print(f"size CSV {cs:.1f} MB vs Parquet {ps:.1f} MB ({cs/ps:.1f}x)")
print(f"one col CSV {tc:.3f}s vs Parquet {tp:.3f}s ({tc/tp:.0f}x)")
print("\nThe CSV must still scan every byte to find one column. Parquet just seeks to it.")
size CSV 29.6 MB vs Parquet 5.2 MB (5.7x) one col CSV 0.062s vs Parquet 0.108s (1x) The CSV must still scan every byte to find one column. Parquet just seeks to it.
Challenge 4 · MapReduce for the busiest hour¶
Same three steps, a different key: count rides per pickup hour, then name the peak.
In [5]:
from collections import defaultdict
pairs = [(r.pickup_hour, 1) for r in df.itertuples()] # MAP
sh = defaultdict(list)
for k, v in pairs: sh[k].append(v) # SHUFFLE
counts = pd.Series({k: sum(v) for k, v in sh.items()}).sort_index() # REDUCE
print(counts.to_string())
print(f"\nbusiest hour: {counts.idxmax()}:00 with {counts.max():,} rides")
print("matches value_counts?", counts.equals(df.pickup_hour.value_counts().sort_index()))
fig, ax = plt.subplots(figsize=(8, 3))
ax.bar(counts.index, counts.values, color=FU); ax.set_xlabel("pickup hour"); ax.set_ylabel("rides")
ax.set_title("Two rush hours, found by MapReduce"); plt.tight_layout(); plt.show()
1 1 2 12 3 45 4 167 5 549 6 1209 7 2086 8 2721 9 2586 10 1812 11 974 12 413 13 345 14 610 15 1126 16 1758 17 2229 18 2130 19 1796 20 1262 21 696 22 323 23 150 busiest hour: 8:00 with 2,721 rides matches value_counts? True
Challenge 5 · Find the skew, then fix it¶
Report how lopsided the partitions are, salt the hot key, and confirm the answer survives.
In [6]:
parts = df.city.value_counts()
imbalance = parts.max() / parts.mean()
print(f"largest partition {parts.max():,} rows, average {parts.mean():,.0f} -> imbalance {imbalance:.1f}x")
for SALT in [1, 2, 4, 8]:
salted = df.city + "#" + np.random.default_rng(0).integers(0, SALT, len(df)).astype(str)
v = salted.value_counts()
print(f" salt={SALT}: {len(v):>2} partitions, largest holds {100*v.max()/len(df):4.1f}% of rows")
salted = df.city + "#" + np.random.default_rng(0).integers(0, 8, len(df)).astype(str)
st1 = df.groupby(salted).fare.sum()
final = st1.groupby([k.split("#")[0] for k in st1.index]).sum()
print("\nrevenue unchanged after salting?", np.allclose(final.sort_index(), df.groupby("city").fare.sum().sort_index()))
print("\nSalting costs a second aggregation stage. It is worth it only when one key really does dominate.")
largest partition 11,531 rows, average 4,167 -> imbalance 2.8x salt=1: 6 partitions, largest holds 46.1% of rows salt=2: 12 partitions, largest holds 23.3% of rows salt=4: 24 partitions, largest holds 11.7% of rows salt=8: 48 partitions, largest holds 5.9% of rows revenue unchanged after salting? True Salting costs a second aggregation stage. It is worth it only when one key really does dominate.