How to recursively call FlatMapFunction<Row,Row>?

Viewed 133

I have a Dataset df, read using spark.read().json

It's schema is something like the following:

root
 |-- items: struct (nullable = true)
 |    |-- item: array (nullable = true)
 |    |    |-- element: struct (containsNull = true)
 |    |    |    |-- batch-id column: long (nullable = true)
 |    |    |    |-- id: string (nullable = true)
 |    |    |    |-- name: string (nullable = true)
 |    |    |    |-- type: string (nullable = true)

I want use a FlatMapFunction to get a dataset with the inner schema(id,name,type).

I want to do something like the following:

df.flatMap(mapperFunction(),RowEncoder.apply(someSchema);

public static FlatMapFunction<Row,Row> mapperFunction() {
    return row -> {
      Row r1 = row.getAs("items");
      List<Row> r2 = r1.getList(0);    //This will explode the column
      StructType schema = r2.get(0).schema();
      //I know list doesn't have map function, I want to know what can be done here
      return r2.flatMap(mapperFunction(),RowEncoder.apply(schema);    
    };
  }
1 Answers

There are several options:

Option 1: use explode

The easiest way to flatten the data structure would be to use explode instead of a flatMap call:

Starting with the data

{"items": {"item": [{"batch-id":1,"id":"id1","name":"name1","type":"type1"},{"batch-id":2,"id":"id2","name":"name2","type":"type2"}]}}

the code

df.withColumn("exploded", explode(col("items.item"))).select("exploded.*").show();

prints

+--------+---+-----+-----+
|batch_id| id| name| type|
+--------+---+-----+-----+
|       1|id1|name1|type1|
|       2|id2|name2|type2|
+--------+---+-----+-----+

Option 2: use flatMap

If the flatMap call is required (for example to add more logic into the mapping), this code prints the same result:

df.flatMap(mapperFunction(), Encoders.bean(Data.class)).show();

with the mapping function

private static FlatMapFunction<Row, Data> mapperFunction() {
    return row -> {
        Row r1 = row.getAs("items");
        List<Row> r2 = r1.getList(0);    //This will explode the column
        return r2.stream().map(entry -> {
            Data d = new Data();
            d.setBatch_id(entry.getLong(0));
            d.setId(entry.getString(1));
            d.setName(entry.getString(2));
            d.setType(entry.getString(3));
            return d;
        }).iterator();
    };
}

and the Data bean

public static class Data implements Serializable {
    private long batch_id;
    private String id;
    private String name;
    private String type;

    //getters and setters
}
Related