One dataset, two engines: scaling California housing until pandas gives up, then rerunning it on pyspark.pandas

Lesson 11 mapped what the pandas API on Spark implements. This lesson measures. We take a dataset every DS course uses — scikit-learn’s California housing, 20,640 rows — and multiply it with jitter until it stops fitting: 5 million rows, 20 million, 51 million, 103 million. At every rung we run the same EDA-and-feature workflow twice: once in plain pandas, once in pyspark.pandas, on the same machine, and record wall-clock time and peak memory for the whole process tree. The pandas side gets a hard 16 GB address-space cap — when it dies, it dies with a MemoryError, exactly the way your notebook kernel dies, and we report the rung where that happens.

Everything here was actually run, and the numbers below are the measured results — not estimates. Hardware and versions, for honesty: a 6-core/12-thread Intel Core i5-10400 (2.9 GHz) with 32 GB RAM and NVMe storage; Python 3.13, pandas 2.3.3, PySpark 4.2.0 running in local[12] mode. No cluster is involved, so the comparison isolates the engine, not the hardware.

Building an “unfittable” dataset from a scikit-learn one

fetch_california_housing gives 20,640 rows and 9 numeric columns (8 features plus the MedHouseVal target) — about 1.5 MB in memory. That is not big data; it is barely data. We scale it by replication with jitter: each replica is the full dataset plus 1%-of-a-standard-deviation Gaussian noise on every numeric column, and ~1 km of noise on the coordinates. Replicas are statistically indistinguishable from the original but not byte-identical — dedup can’t cheat, compression can’t cheat, and groupby cardinality behaves realistically.

import numpy as np
import pandas as pd
from sklearn.datasets import fetch_california_housing

base = fetch_california_housing(as_frame=True).frame     # 20,640 rows x 9 cols
num_cols = [c for c in base.columns if c not in ("Latitude", "Longitude")]
stds = base[num_cols].std()

rng = np.random.default_rng(42)
def replica() -> pd.DataFrame:
    rep = base.copy()
    for c in num_cols:
        rep[c] = rep[c] + rng.normal(0.0, stds[c] * 0.01, len(rep))
    rep["Latitude"]  = rep["Latitude"]  + rng.normal(0.0, 0.01, len(rep))
    rep["Longitude"] = rep["Longitude"] + rng.normal(0.0, 0.01, len(rep))
    return rep

# 20 parquet files x 250 replicas each = 5,000x = 103.2M rows on disk
for chunk in range(20):
    pd.concat([replica() for _ in range(250)], ignore_index=True) \
      .to_parquet(f"data/part-{chunk:02d}.parquet", index=False)

Writing the ladder as 20 equal parquet files (~7.5 GB total) gives us every rung for free: 1 file is 250× (5.16M rows), 4 files are 1,000× (20.6M), 10 files are 2,500× (51.6M), all 20 are 5,000× (103.2M rows). The original dataset is kept as its own file for the 1× rung — that rung exists to show the other side of the trade, where Spark loses.

The workflow under test

A deliberately ordinary mid-notebook sequence — load, engineer features, aggregate to a spatial grid, join the aggregates back, summarize:

# 1. load
df = pd.concat([pd.read_parquet(f) for f in files], ignore_index=True)

# 2. feature engineering
df["rooms_per_person"] = df["AveRooms"] / df["AveOccup"]
df["bedrm_ratio"]      = df["AveBedrms"] / df["AveRooms"]
df["log_pop"]          = np.log1p(df["Population"])
df["lat_cell"]         = df["Latitude"].round(1)
df["lon_cell"]         = df["Longitude"].round(1)

# 3. aggregate to a 0.1-degree grid
cells = (df.groupby(["lat_cell", "lon_cell"])
           .agg({"MedHouseVal": "mean", "MedInc": "mean", "HouseAge": "count"})
           .rename(columns={"MedHouseVal": "cell_value",
                            "MedInc": "cell_income", "HouseAge": "n"})
           .reset_index())

# 4. join the aggregates back; relative-value feature
df = df.merge(cells, on=["lat_cell", "lon_cell"], how="left")
df["rel_value"] = df["MedHouseVal"] / df["cell_value"]

# 5. small outputs: top cells, a mean, a correlation
top = cells.sort_values("n", ascending=False).head(5)
mean_rel = df["rel_value"].mean()
corr = df["MedInc"].corr(df["MedHouseVal"])

Every step is textbook pandas. Step 4 is the memory villain: merge materializes a second copy of the frame.

And here is the entire diff for the pyspark.pandas version:

import pyspark.pandas as ps
ps.set_option("compute.default_index_type", "distributed")

df = ps.read_parquet("data/part-{00,01,02,03}.parquet")   # replaces pd.concat([...])

That is it. Steps 2–5 run character-for-character identical — the arithmetic, the np.log1p ufunc, the dict-agg groupby, the merge, the correlation. (The only other touch: the tiny top frame is collected with .to_pandas(), since we want it on the driver.) This is the entire pitch of the pandas API on Spark, and it survives contact with real code.

Results

Workflow time is the measured in-process time; the pyspark.pandas side additionally pays ~5 s of JVM/session startup per script run, which is listed separately rather than hidden.

Scale	Rows	pandas	pandas peak RSS	pyspark.pandas	ps peak RSS
1×	20,640	0.08 s	0.19 GB	5.8 s	0.9 GB
250×	5.16 M	4.9 s	2.8 GB	13.2 s	2.9 GB
1,000×	20.64 M	18.4 s	11.3 GB	34.2 s	5.0 GB
2,500×	51.6 M	💀 MemoryError at 69 s (in merge)	11.4 GB at death	62.6 s	5.9 GB
5,000×	103.2 M	💀 MemoryError at 32 s (in load)	—	113.0 s	5.8 GB

Reading the table honestly:

pandas is unbeatable while it fits. At 1× the workflow takes 80 milliseconds in pandas and 5.8 s in Spark — a 70× penalty for distributing 20 thousand rows, before counting JVM startup. Even at 5 million rows pandas is 2.7× faster; at 20 million, still 1.9× faster. This is lesson 10’s small-data trap measured from the other side. On a single machine, Spark in local mode never out-ran live pandas in this experiment — not at any rung where pandas survived.

What Spark won was survival, and that is the win that matters. At 51.6M rows the comparison stops being “2× slower” and becomes “only one of them finishes”. pandas died 69 seconds in, inside merge‘s join-indexer, with the allocator refusing a 394 MiB array on top of an 11.4 GB process. At 103.2M rows it died earlier and harder: the load itself — pd.concat assembling the final frame — asked for one 6.9 GB contiguous block and was refused. The cap converts “my laptop froze and I rebooted it” into an honest, reportable error; without it, the same workload swaps the machine into paralysis. Meanwhile pyspark.pandas finished the 103M-row rung in 113 s, and — this is the key column — its peak memory stayed flat around 5–6 GB from 20M to 103M rows, because executors stream partitions and spill shuffles to disk instead of holding everything at once. pandas’ footprint scales linearly with rows; Spark’s footprint scales with partition size. That flatness is what “scales horizontally” actually looks like from a process monitor.

Would more RAM have saved pandas? At 2,500×, yes: uncapped, the merge needs roughly 28 GB, so a 64 GB workstation clears it. That is a legitimate answer — until the next multiplier. The wall doesn’t disappear with bigger hardware; it moves one rung.

Lazy evaluation shows up in the stage timings. Per-stage timers in the Spark runs are almost comic: at 103M rows, “load” reports 1.4 s, “features” 0.2 s, “merge” 0.2 s — and then “summarize” reports 110.7 s, because that’s where the first real action sits and the whole accumulated plan executes. If you instrument a pyspark.pandas notebook the way you instrument pandas, you will conclude that computing a mean is your bottleneck. It isn’t; it’s merely where the bill arrives (lesson 11’s execution-model section, demonstrated).

The index-type tax, measured — and a surprise

Lesson 11 warned that the manufactured default index (compute.default_index_type) is a hidden cost. So we ran the whole 103M-row workflow under each setting, expecting distributed-sequence to pay an extra pass over the data. It didn’t: 113.0 s vs 114.6 s — identical within run-to-run noise. Even inserting an explicit df = df.reset_index() after the read changed nothing, and the notionally catastrophic sequence type clocked the same 113 s.

The explanation is the best kind of surprise: Spark’s optimizer pruned the index away. Our workflow’s groupbys, merges, and column arithmetic never consume row identity, so the Catalyst column-pruner removes the index computation from the physical plan entirely — even after reset_index(), because the resulting column is never read. The tax is not merely lazy; it is elided when unused.

So we forced the issue: build the index and actually consume it (df.reset_index()["index"].max()) on all 103.2M rows. Now the three strategies separate exactly as lesson 11’s table predicts:

compute.default_index_type	Build + consume the index, 103.2M rows	Max id produced
distributed	2.6 s	506,808,155,199 — non-continuous, huge gaps
distributed-sequence (default)	15.2 s	103,199,999 — continuous 0…N−1
sequence	68.4 s — final stage runs on one partition	103,199,999

Three practical readings. First, the distributed ids are real: the max id for 103 million rows is 506 billion — anything downstream assuming dense 0…N−1 ids breaks. Second, the default costs ~6× the fast option when the index is consumed — and nothing when it isn’t, so don’t cargo-cult distributed into every job; set it when your pipeline actually touches row identity. Third, sequence announces its pathology right in the Spark UI: a job over 60 partitions ends in a stage of exactly one task. If you ever see that shape, you now know what it is.

Validating that both engines agree

A benchmark where the engines compute different answers is a bug report, not a benchmark. On every jittered rung, the two implementations agreed exactly on the structure — same row counts, same number of grid cells (2,285 / 2,374 / 2,432 / 2,476 as the jitter widens the grid), identical top cell — and agreed on the numerics to floating-point-reduction precision: mean_rel_value matches to 13 decimal places, the income–value correlation to 11 (pandas 0.68800797057857 28 vs Spark ...613 at 20.6M rows). The residual is reduction order: distributed sums add the same numbers in a different order. “Close to machine precision, not bit-identical” is the correct expectation when porting, and worth encoding in your tests as tolerances rather than equality.

The un-jittered 1× rung produced the one genuine divergence, and it is instructive: pandas counted 1,563 grid cells, Spark 1,577 — from character-identical code. The culprit is round(1): pandas/NumPy round half-to-even (banker’s rounding), Spark SQL rounds half-up, and the raw dataset’s coordinates are quoted to two decimals, so values like 34.05 sit exactly on a cell boundary and get rounded into different cells by the two engines. The jittered rungs are immune (noise makes exact .x5 boundaries measure-zero). Two lessons in one: same-API-different-rounding is precisely the “same name, different behavior” bucket from lesson 11 — and any binning scheme that puts real data exactly on bin edges is fragile in any engine; bin with np.floor(x * 10) / 10 if you need cross-engine determinism.

Where the model fits in

This track is about ML, so: where does the estimator go once features are distributed? Three patterns, in order of preference:

1. Aggregate, then sklearn. If the model trains on aggregates or a sample — the common case — do the heavy lifting distributed, then sample(frac=...).to_pandas() and hand scikit-learn what fits. The 103M-row frame above collapses to 2,476 grid cells; any laptop trains on that.
2. Grouped training. One model per segment (per grid cell, per region): applyInPandas from lesson 8 — pandas code runs on executors, one group at a time.
3. Spark MLlib. When one global model must see every row (lessons 3–7). pyspark.pandas frames hand off via .to_spark() and a VectorAssembler.

What you should not do is df.to_pandas() on the full frame to “just use sklearn” — that is the same collect that kills notebooks, and the pandas column of the results table shows exactly how much driver memory it would need.

Take-aways

• Replication-with-jitter turns any toy dataset into an honest scaling ladder; 20 parquet files give you every rung for free.
• The port was one import plus one changed read call; the workflow body did not change by a character.
• pandas won every rung it survived (70× at 20k rows, ~2× at 20M) and lost by dying: MemoryError in the merge at 51.6M rows, in the load at 103M. Spark’s memory stayed flat at 5–6 GB throughout.
• Per-stage timing lies in a lazy engine; the action pays for the whole plan.
• Identical code can still diverge on edge semantics — banker’s rounding vs half-up put 14 grid cells on different sides of a boundary. Validate ports with tolerances, and don’t bin on exact half-way values.
• Keep collects at the edges: aggregates out, samples out, never the full frame.

Next lesson closes the trilogy: interop between the three frame types, the recurring performance traps and their fixes, and a checklist for migrating an existing pandas codebase.