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.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.ArrayType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StringType;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import java.util.Arrays;
import java.util.List;

public class FlatMapExample {
    public static void main(String[] args) {
        SparkSession spark = SparkSession.builder()
            .appName("Spark flatMap example")
            .master("local[*]")
            .getOrCreate();

        List<Row> arrayStructureData = Arrays.asList(
            RowFactory.create("James,,Smith", Arrays.asList("Java", "Scala", "C++"), "CA"),
            RowFactory.create("Michael,Rose,", Arrays.asList("Spark", "Java", "C++"), "NJ"),
            RowFactory.create("Robert,,Williams", Arrays.asList("CSharp", "VB", "R"), "NV")
        );

        StructType arrayStructureSchema = new StructType(new StructField[] {
            new StructField("name", DataTypes.StringType, false, null),
            new StructField("languagesAtSchool", new ArrayType(DataTypes.StringType, false), false, null),
            new StructField("currentState", DataTypes.StringType, false, null)
        });

        Dataset<Row> df = spark.createDataFrame(arrayStructureData, arrayStructureSchema);

        Dataset<Row> df2 = df.flatMap(row -> {
            List<String> languages = row.getList(1);
            return languages.stream()
                .map(language -> RowFactory.create(row.getString(0), language, row.getString(2)))
                .iterator();
        }, RowEncoder.apply(new StructType(new StructField[] {
            new StructField("Name", StringType, false, null),
            new StructField("Language", StringType, false, null),
            new StructField("State", StringType, false, null)
        })));

        df2.show(false);

        spark.stop();
    }
}

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   |
+-----------------+---------+-----+