Introduction

Spark flatMap() transformation flattens the DataFrame column after applying the function on every element and returns a new DataFrame respectively.

The returned DataFrame can have the same count or more number of elements. This is one of the major differences between flatMap() and map(), where map() transformation always returns the same number of elements as input.

import org.apache.spark.sql.{Row, SparkSession}
import org.apache.spark.sql.types.{ArrayType, StringType, StructType}
import org.apache.spark.sql.functions.col

val arrayStructureData = Seq(
    Row("James,,Smith", List("Java", "Scala", "C++"), "CA"),
    Row("Michael,Rose,", List("Spark", "Java", "C++"), "NJ"),
    Row("Robert,,Williams", List("CSharp", "VB", "R"), "NV")
)

val arrayStructureSchema = new StructType()
    .add("name", StringType)
    .add("languagesAtSchool", ArrayType(StringType))
    .add("currentState", StringType)

val df = spark.createDataFrame(spark.sparkContext.parallelize(arrayStructureData), arrayStructureSchema)
import spark.implicits._

// flatMap() Usage
val df2 = df.flatMap(f => f.getSeq[String](1).map((f.getString(0), _, f.getString(2))))
    .toDF("Name", "Language", "State")

df2.show(false)

This yields the following output after the flatMap() transformation. As you notice, the input of the DataFrame has 3 records, but after exploding the "language" using flatMap(), it returns 9 elements.

+-----------------+---------+-----+
|Name             |Language |State|
+-----------------+---------+-----+
|James,,Smith     |Java     |CA   |
|James,,Smith     |Scala    |CA   |
|James,,Smith     |C++      |CA   |
|Michael,Rose,    |Spark    |NJ   |
|Michael,Rose,    |Java     |NJ   |
|Michael,Rose,    |C++      |NJ   |
|Robert,,Williams |CSharp   |NV   |
|Robert,,Williams |VB       |NV   |
|Robert,,Williams |R        |NV   |
+-----------------+---------+-----+