Living with pyspark.pandas: interop, performance traps, and a pandas-to-Spark migration checklist

Lesson 11 mapped what the pandas API on Spark implements; lesson 12 proved on 103 million rows that the port can be a handful of changed lines. This closing lesson is about operating it: moving data between the three frame types without accidental collects, the recurring performance traps and their fixes, the escape hatches when the pandas layer isn’t the right tool for a step — and a migration checklist you can run against an existing pandas codebase.

Three frame types, six conversions

Production jobs rarely live in one API. You will hold data as a plain pandas DataFrame (driver-local), a pyspark.pandas frame (distributed, pandas-shaped), and a native Spark DataFrame (distributed, Spark-shaped) — often all three in one pipeline.

import pandas as pd
import pyspark.pandas as ps

pdf = pd.read_parquet("sample.parquet")   # driver-local pandas

psdf = ps.from_pandas(pdf)                # pandas  -> pyspark.pandas  (distributes the data)
sdf  = psdf.to_spark(index_col="id")      # ps      -> native Spark    (metadata-only, no copy)
psdf = sdf.pandas_api(index_col="id")     # Spark   -> pyspark.pandas  (metadata-only, no copy)
pdf  = psdf.to_pandas()                   # ps      -> pandas          (COLLECTS to the driver)

The cost model matters more than the syntax:

Conversion	Data movement
ps.from_pandas(pdf)	driver → cluster upload; fine for lookup tables, wrong for big data (read from storage instead)
.to_spark() / .pandas_api()	free — same distributed data, different API wrapper
.to_pandas(), .to_numpy(), .values	full collect to driver memory — the same toPandas() trap from lesson 2, wearing pandas clothes

Two habits keep you safe. First, thread index_col through to_spark()/pandas_api(); if you don’t, the index is either silently dropped (one direction) or re-manufactured (the other), and you re-buy the default-index cost from lesson 11 on every crossing. Second, grep your code for .to_pandas() and make each occurrence justify itself with a size argument — “it’s aggregated to 2 thousand rows” is a justification; “the next library call wanted pandas” starts a driver-OOM investigation.

The collect is legitimate exactly when distributed work is done: aggregates going into a plot, a model-ready sample, a small scored output. That is the sample-to-pandas hand-off pattern from lesson 12, and it pairs with the reverse hand-off — applyInPandas/apply_batch from lesson 8 — when you want pandas code to run on the cluster instead of pandas data on the driver.

The recurring performance traps

Each of these has appeared once in this track; here they are as a single checklist with fixes.

Trap 1 — the manufactured index. Any operation that needs row identity on a frame without a real index triggers default-index construction; with the distributed-sequence default that is an extra pass over the data before your actual work. Fix: index_col= at every read and every to_spark()/pandas_api() crossing; or ps.set_option("compute.default_index_type", "distributed") when you never rely on positional semantics. Lesson 12 measured the difference on the 103M-row workflow: the extra pass is real money on a real cluster.

Trap 2 — sort_values as a display habit. In pandas, sorting is cheap enough to be a reading aid. Distributed, it is a full shuffle, and a .sort_values(...).head(10) runs that shuffle to answer a top-10 question nlargest(10, "col") answers with a cheap parallel reduction. Fix: sort only when output order is a requirement; use nlargest/nsmallest for top-k.

Trap 3 — row-wise apply with an unannotated function. df.apply(f, axis=1) ships batches through Arrow into Python — orders of magnitude slower than column arithmetic that compiles to Spark SQL expressions, and schema inference runs f on a sample first. Fix: rewrite as column expressions when possible; when not, annotate the return type and prefer apply_batch/transform_batch, which hand your function a whole pandas chunk:

def enrich(chunk: pd.DataFrame) -> pd.DataFrame:      # runs on executors, real pandas inside
    chunk["geohash"] = encode(chunk["lat"], chunk["lon"])
    return chunk

df = df.pandas_on_spark.apply_batch(enrich)

Trap 4 — the ever-growing lazy plan. A feature-factory loop that assigns 300 columns builds a 300-step logical plan; optimization time itself becomes the bottleneck, and a failure at the end recomputes everything. Fix: df = df.spark.cache() after expensive stable stages (remember to df.spark.unpersist()), or df = df.spark.local_checkpoint() to cut lineage entirely. Rule of thumb from lesson 10 applies unchanged: checkpoint where you’d have written an intermediate table.

Trap 5 — cross-frame alignment. compute.ops_on_diff_frames turns index alignment between different frames into silent joins — and on Spark 4.2 it ships enabled by default (lesson 11). One innocent-looking df["n"] = other["count"] is a shuffle. Fix: ps.set_option("compute.ops_on_diff_frames", False) at session start and write explicit merges; when you genuinely want an aligned op, enable it inside ps.option_context(...) so the blast radius is one block.

Trap 6 — collect-shaped conveniences. transpose(), .values, for row in df.iterrows(), df.to_numpy() — all either guarded by compute.max_rows or unguarded collects. The guard default (1000 rows) is there to be respected, not raised on reflex.

Escape hatches: dropping below the pandas layer

The pandas layer is a view over a Spark DataFrame — use that. Three hatches, cheapest first.

Spark SQL over a pandas-shaped frame. ps.sql interpolates frames directly and returns a ps.DataFrame:

top = ps.sql("""
    SELECT lat_cell, lon_cell, avg(MedHouseVal) AS cell_value, count(*) AS n
    FROM {df}
    GROUP BY lat_cell, lon_cell
    HAVING count(*) > 100
""", df=df)

Native DataFrame API for a step. Window functions, Bucketizer-style quantile bucketing (the qcut gap from lesson 11), broadcast hints, or fine repartitioning are one .to_spark() away — do the step, come back with .pandas_api(), keep index_col threaded through both.

Spark MLlib when the model itself must be distributed. That is lessons 3–7; the boundary rule from lesson 12 stands: features distributed, then either sample-to-pandas for sklearn or stay in Spark for MLlib.

The skill being practiced is knowing which layer a step belongs to. A pipeline that is 90% pandas-shaped with two surgical .to_spark() steps is idiomatic, not impure.

Migrating an existing pandas codebase: the checklist

Run this against a real repository before committing to a port.

Phase 0 — decide whether to port at all.
– Data actually bigger than one machine’s memory, or growing there? If a cx52-sized box holds it, lesson 10’s small-data trap says stay in pandas (or reach for Polars — see the earlier Polars-vs-pandas piece on this blog).
– Is the slow part IO or compute? Distributed reads fix IO-bound jobs; a single quadratic pandas step needs an algorithm fix, not a cluster.

Phase 1 — inventory the incompatibilities. Grep for the patterns that do not port:
– .iterrows() / .itertuples() / .iat / .at loops → rewrite as column ops or apply_batch before the port; they are also your worst pandas code today.
– .values / .to_numpy() mid-pipeline → these become collects; replace with column ops.
– qcut / exact quantile dependencies → approximate quantiles or Bucketizer.
– Implicit row-order assumptions (head() after concat, positional slicing) → add explicit sorts or keys now.
– inplace=True everywhere → harmless but rewrite to assignment; it clarifies the lazy semantics.
– Chained indexing (df[a][b] = x) → broken in modern pandas anyway; fix first.

Phase 2 — port the IO boundary first. Move reads to ps.read_parquet(..., index_col=...) and writes to df.to_parquet(...); keep every transformation as-is; run the suite. Most failures at this stage are the four buckets from lesson 11, and each error message names the offending method.

Phase 3 — de-trap. Apply the six traps above; add spark.cache() at stage boundaries; kill collect-shaped conveniences.

Phase 4 — validate numerically, not just structurally. Approximate quantile/corr, floating-point reduction order, and null-handling differences mean outputs are close, not identical. Compare aggregates with tolerances (the lesson 12 workflow agreed to 11–13 decimal places on its summary statistics — but diverged on grid-cell counts through a rounding-rule difference); pin down any column where “close” is not acceptable and compute it exactly via SQL.

Phase 5 — right-size. After the port works, profile: df.spark.explain() shows the actual plan behind any frame — the habit of reading plans from lesson 2 transfers directly. Confirm partition counts are sane after big joins, and that the job is not secretly one sequence-indexed stage.

When not to use it

An honest tool description includes its non-uses:

• Data fits in memory — plain pandas or Polars wins on latency, debuggability, and ecosystem (lesson 10, measured again in lesson 12’s 1× rung).
• The workload is one big custom algorithm, not dataframe transformations — go straight to native Spark or applyInPandas; a pandas facade adds nothing.
• You need exact pandas semantics — regulatory replays, byte-identical backtests. The approximate and order-sensitive corners make “close” the contract, and close may not be good enough.
• Streaming — the pandas API is batch; Structured Streaming is a different course.

Take-aways

The pandas API on Spark earns its place when you treat it as what it is: a distributed engine with a familiar steering wheel. Thread real indexes through every boundary, keep collects at the edges and justified, batch your Python through Arrow instead of dripping it row by row, checkpoint long lineages, and drop to SQL or the native API for the steps that want it. Ported that way, the pandas notebook that died at 100 million rows in lesson 12 becomes a job that scales with the cluster — and the checklist above turns “we should move this to Spark someday” into a bounded engineering task rather than a rewrite.

This closes the Distributed ML with PySpark track’s pandas trilogy: what the API is (lesson 11), what it costs and saves in measured numbers (lesson 12), and how to live with it in production (this lesson).