Introduction

PySpark collect() and collectAsList() are action operations used to bring all elements from a DataFrame or Dataset (across all cluster nodes) back to the driver in PySpark. These operations should be used only on small datasets, typically after applying transformations like filter(), groupBy(), or count(). Using collect() on large datasets can cause driver out-of-memory errors.

In order to explain with an example, first, let’s create a DataFrame.

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, Row

spark = SparkSession.builder \
    .master("local[1]") \
    .appName("AIStudyLabs.com") \
    .getOrCreate()

data = [
    Row(name=Row(firstname="James", middlename="", lastname="Smith"), id="36636", gender="M", salary=3000),
    Row(name=Row(firstname="Michael", middlename="Rose", lastname=""), id="40288", gender="M", salary=4000),
    Row(name=Row(firstname="Robert", middlename="", lastname="Williams"), id="42114", gender="M", salary=4000),
    Row(name=Row(firstname="Maria", middlename="Anne", lastname="Jones"), id="39192", gender="F", salary=4000),
    Row(name=Row(firstname="Jen", middlename="Mary", lastname="Brown"), id="", gender="F", salary=-1)
]

schema = StructType([
    StructField('name', StructType([
        StructField('firstname', StringType(), True),
        StructField('middlename', StringType(), True),
        StructField('lastname', StringType(), True)
    ])),
    StructField('id', StringType(), True),
    StructField('gender', StringType(), True),
    StructField('salary', IntegerType(), True)
])

df = spark.createDataFrame(data, schema)
df.printSchema()
df.show(truncate=False)

root
 |-- name: struct (nullable = true)
 |    |-- firstname: string (nullable = true)
 |    |-- middlename: string (nullable = true)
 |    |-- lastname: string (nullable = true)
 |-- id: string (nullable = true)
 |-- gender: string (nullable = true)
 |-- salary: integer (nullable = true)

+---------------------+-----+------+------+
|name                 |id   |gender|salary|
+---------------------+-----+------+------+
|{James, , Smith}     |36636|M     |3000  |
|{Michael, Rose, }    |40288|M     |4000  |
|{Robert, , Williams} |42114|M     |4000  |
|{Maria, Anne, Jones} |39192|F     |4000  |
|{Jen, Mary, Brown}   |     |F     |-1    |
+---------------------+-----+------+------+

Using collect() and collectAsList()

The collect() action function is used to retrieve all elements from the DataFrame as a list of Row objects to the driver program.

collectAsList() is not a method in PySpark, so in Python, we typically use collect() to achieve the same outcome.

Example

col_list = df.collect()

for row in col_list:
    salary = row['salary']
    print(salary)

3000
4000
4000
4000
-1

Retrieving data from Struct column

To retrieve a struct column from Row, we should access the fields directly.

# Retrieving data from Struct column
for row in col_list:
    salary = row['salary']
    fullName = row['name']
    firstName = fullName['firstname']
    middleName = fullName['middlename']
    lastName = fullName['lastname']
    print(f"{firstName},{middleName},{lastName},{salary}")

James,,Smith,3000
Michael,Rose,,4000
Robert,,Williams,4000
Maria,Anne,Jones,4000
Jen,Mary,Brown,-1

Note that like other DataFrame functions, collect() does not return a DataFrame; instead, it returns data in a list to your driver. Once the data is collected in a list, you can use Python's native capabilities for further processing.

In case you want to just return certain elements of a DataFrame, you should call select() first.

data_collect = df.select("name").collect()

When to avoid collect()

Usually, collect() is used to retrieve the action output when you have a very small result set, and calling collect() on a DataFrame with a larger result set causes out-of-memory errors as it returns the entire dataset (from all workers) to the driver. Hence, we should avoid calling collect() on a larger dataset.

collect() vs select()

The select() method on a DataFrame returns a new DataFrame that holds the columns that are selected, whereas collect() returns the entire dataset to the driver.

select() is a transformation function, whereas collect() is an action.