Skip to content
LakeBench
ProblemsCommunityPricing
Sign inStart practicing
Back
  1. Home
  2. Interview prep
  3. toPandas and collect safely

PySpark · DataFrame API in Practice

toPandas and collect safely

Easypyspark-82
collecttopandasdriverarrow

Question

When is it safe to call toPandas() or collect()?

Solution

Only when the result is small. collect() and toPandas() bring every row to the driver's memory, so they are safe on a few thousand or a few hundred thousand rows, not on a table.

What happens

Executors send their partitions to the driver, which holds all rows in one Python process. If the result is bigger than the driver's memory, the driver fails with an out-of-memory error. If it is bigger than spark.driver.maxResultSize (1 GB by default), Spark aborts the job with an error that the serialized results are larger than that limit. The driver is a single machine, so everything the cluster did in parallel ends in one place.

Safe patterns

# aggregate first, then convert
summary = df.groupBy("region").agg(F.sum("amount").alias("revenue"))
pdf = summary.toPandas()          # a handful of rows

# for debugging, take a small piece
df.limit(1000).toPandas()
df.sample(0.001, seed=1).toPandas()
df.show(20)

Speeding up toPandas

Converting through Python row objects is slow. Apache Arrow transfers data in columnar batches and converts much faster. Turn it on with spark.sql.execution.arrow.pyspark.enabled=true (check the default in your Spark version, since it differs). Arrow also supports Pandas UDFs and createDataFrame(pandas_df).

Large outputs

If the result is large, do not bring it to the driver. Write it to storage (df.write.parquet(...)) and let the next step read it from there. If you need to process each partition in Python, use mapInPandas or foreachPartition, which run on the executors.

A typical mistake

Someone calls toPandas() on a 200 million row DataFrame to plot it. The notebook hangs, then the driver dies. The fix is to aggregate or sample before plotting. Another typical one is calling collect() to loop in Python over rows and perform per-row work, which discards Spark's parallelism. Use DataFrame operations or foreachPartition instead.

🎯 Put this concept into practice

Solidify this answer with real hands-on interview drills in the browser studio.

Open related drill →
PreviousNext