Introduction

Spark collect() and collectAsList() are action operations used to bring all elements from an RDD, DataFrame, or Dataset (across all cluster nodes) back to the driver. These operations should be used only on small datasets, typically after applying transformations like filter(), group(), 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.

// Add imports here from what you have learned so far
SparkSession spark = SparkSession.builder()
    .master("local[1]")
    .appName("AIStudyLabs.com")
    .getOrCreate();

List<Row> data = Arrays.asList(
    RowFactory.create(RowFactory.create("James ", "", "Smith"), "36636", "M", 3000),
    RowFactory.create(RowFactory.create("Michael ", "Rose", ""), "40288", "M", 4000),
    RowFactory.create(RowFactory.create("Robert ", "", "Williams"), "42114", "M", 4000),
    RowFactory.create(RowFactory.create("Maria ", "Anne", "Jones"), "39192", "F", 4000),
    RowFactory.create(RowFactory.create("Jen", "Mary", "Brown"), "", "F", -1)
);

StructType schema = new StructType()
    .add("name", new StructType()
        .add("firstname", DataTypes.StringType)
        .add("middlename", DataTypes.StringType)
        .add("lastname", DataTypes.StringType))
    .add("id", DataTypes.StringType)
    .add("gender", DataTypes.StringType)
    .add("salary", DataTypes.IntegerType);

Dataset<Row> df = spark.createDataFrame(data, schema);
df.printSchema();
df.show(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()

collect() action function is used to retrieve all elements from the dataset (RDD/DataFrame/Dataset) as an array to the driver program.

collectAsList() action function is similar to collect() but it returns a Java util list.

Example

List<Row> colList = df.collectAsList();
Row[] colData = (Row[]) df.collect();

for (Row row : colData) {
    int salary = row.getInt(3); // Index starts from zero
    System.out.println(salary);
}

3000
4000
4000
4000
-1

Retrieving data from Struct column

To retrieve a struct column from Row, we should use getStruct() function.

// Retrieving data from Struct column
for (Row row : colData) {
    int salary = row.getInt(3);
    Row fullName = row.getStruct(0); // Index starts from zero
    String firstName = fullName.getString(0); // In struct row, again index starts from zero
    String middleName = fullName.getString(1);
    String lastName = fullName.getString(2);
    System.out.println(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 an array to your driver. Once the data is collected in an array, you can use Java language for further processing.

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

Row[] dataCollect = (Row[]) 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 an RDD/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 an RDD/DataFrame returns a new DataFrame that holds the columns that are selected, whereas collect() returns the entire dataset.

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

This example is also available at Spark Github project.