The collect() Trap
collect() runs the chain and then sends every result row from every executor to the driver, where the rows become Python objects in a list. The executors hold the data spread across many machines, and the driver is a single process, sized for coordinating tasks rather than for holding data. A collect() on a result that fits on one machine is fine. A collect() on a result that doesn't is the most common way a beginner takes down a driver, and it fails only after every executor has done its work.
- 1g
The funnel
Take 100 executors finishing a job that produces 50 million rows of 120 bytes. Across the cluster that is 6GB, spread 60MB per executor, which each of them held comfortably. collect() asks all 100 to send their 60MB to the driver, and the driver guards itself with spark.driver.maxResultSize. Once the results that have arrived pass that limit, Spark aborts the job with a SparkException saying the total size of serialized results is bigger than spark.driver.maxResultSize. The executors' work is thrown away, and the failure comes late, after the job has done nearly all its reading. Raising the limit to let 6GB through moves the failure onto the driver itself: its heap is small by default too, it has to hold the whole result at once as Python objects that take more space than the bytes on the wire, and it dies with an OutOfMemoryError. Whether collect() is safe depends only on the size of the result, however large the source is.
| What you want | Use | Why it is safe |
|---|---|---|
| A look at some rows | `show(20)` or `take(5)` | A fixed number of rows reaches the driver |
| A number | `count()`, or `agg(...)` then `collect()` | The aggregate is 1 row before it travels |
| The top 10 by a column | `orderBy(...).limit(10).collect()` | The limit caps what the driver receives |
| All the rows, saved | `write.parquet(path)` | Executors write their own partitions in parallel, the driver receives nothing |
| All the rows, in a Python list | `collect()` only when the count is known to be small | A lookup table of 300 rows fits anywhere |
| A local pandas table | `toPandas()` after an aggregate or a limit | `toPandas()` collects, with the same funnel |
When collect() is the right call
collect() is safe after the chain has already shrunk the data. A groupBy("country").agg(...) over 2 billion events returns at most 1 row per country, a few hundred rows, and collecting them into a Python list to build a report costs nothing. A limit(1000) before collect() caps the list at 1000 rows whatever the source holds. The question to ask before typing it is the row count of the DataFrame on the left of the dot. If you cannot bound that count, write the result instead and read it back from where it landed.
- Do
- Aggregate or limit before you collect
- Use show() or take() for a look at the data
- Write large results to storage and let the executors do it
- Bound the row count of the DataFrame before the dot
- Avoid
- collect() a DataFrame whose row count you cannot state
- Treat toPandas() as free, it collects
- Collect to count, since count() returns 1 number
- Raise driver memory as the first fix for a failed collect()
Which of these lines is safe to run against a 3-billion-row events table on a driver with default settings?
- events.select("user_id", "country").collect(), because 2 columns make each row small.
- events.groupBy("country").agg(F.count("*")).collect(), because the result has 1 row per country.
- events.filter(F.col("country") == "US").toPandas(), because the filter keeps 1 country.
- events.distinct().collect(), because distinct() removes the duplicated rows first.
The cell below collects, and it is safe because limit(5) caps the result at 5 rows before they travel. Run it, then change 5 to 50 and note that the driver still receives a bounded list.
About this lesson
This section is part of the Lazy Evaluation: Beginner lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.