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
val spark: SparkSession = SparkSession.builder()
  .master("local[1]")
  .appName("AIStudyLabs.com")
  .getOrCreate()

val data = Seq(
  Row(Row("James ", "", "Smith"), "36636", "M", 3000),
  Row(Row("Michael ", "Rose", ""), "40288", "M", 4000),
  Row(Row("Robert ", "", "Williams"), "42114", "M", 4000),
  Row(Row("Maria ", "Anne", "Jones"), "39192", "F", 4000),
  Row(Row("Jen", "Mary", "Brown"), "", "F", -1)
)

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

val df = spark.createDataFrame(spark.sparkContext.parallelize(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  Array[Row] to the driver program.

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

Example

val colList = df.collectAsList()
val colData = df.collect()

colData.foreach(row => {
  val salary = row.getInt(3) // Index starts from zero
  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
colData.foreach(row => {
  val salary = row.getInt(3)
  val fullName: Row = row.getStruct(0) // Index starts from zero
  val firstName = fullName.getString(0) // In struct row, again index starts from zero
  val middleName = fullName.get(1).toString
  val lastName = fullName.getAs[String]("lastname")
  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 Scala language for further processing.

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

val dataCollect = 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.