Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to use Spark SQL DataFrame with flatMap?

I am using the Spark Scala API. I have a Spark SQL DataFrame (read from an Avro file) with the following schema:

root
|-- ids: array (nullable = true)
|    |-- element: map (containsNull = true)
|    |    |-- key: integer
|    |    |-- value: string (valueContainsNull = true)
|-- match: array (nullable = true)
|    |-- element: integer (containsNull = true)

Essentially 2 columns [ ids: List[Map[Int, String]], match: List[Int] ]. Sample data that looks like:

[List(Map(1 -> a), Map(2 -> b), Map(3 -> c), Map(4 -> d)),List(0, 0, 1, 0)]
[List(Map(5 -> c), Map(6 -> a), Map(7 -> e), Map(8 -> d)),List(1, 0, 1, 0)]
...

What I would like to do is flatMap() each row to produce 3 columns [id, property, match]. Using the above 2 rows as the input data we would get:

[1,a,0]
[2,b,0]
[3,c,1]
[4,d,0]
[5,c,1]
[6,a,0]
[7,e,1]
[8,d,0]
...

and then groupBy the String property (ex: a, b, ...) to produce count("property") and sum("match"):

 a    2    0
 b    1    0
 c    2    2
 d    2    0
 e    1    1

I would want to do something like:

val result = myDataFrame.select("ids","match").flatMap( 
    (row: Row) => row.getList[Map[Int,String]](1).toArray() )
result.groupBy("property").agg(Map(
    "property" -> "count",
    "match" -> "sum" ) )

The problem is that the flatMap converts DataFrame to RDD. Is there a good way to do a flatMap type operation followed by groupBy using DataFrames?

like image 432
Yuri Brovman Avatar asked May 21 '15 18:05

Yuri Brovman


1 Answers

What does flatMap do that you want? It converts each input row into 0 or more rows. It can filter them out, or it can add new ones. In SQL to get the same functionality you use join. Can you do what you want to do with a join?

Alternatively, you could also look at Dataframe.explode, which is just a specific kind of join (you can easily craft your own explode by joining a DataFrame to a UDF). explode takes a single column as input and lets you split it or convert it into multiple values and then join the original row back onto the new rows. So:

user      groups
griffin   mkt,it,admin

Could become:

user      group
griffin   mkt
griffin   it
griffin   admin

So I would say take a look at DataFrame.explode and if that doesn't get you there easily, try joins with UDFs.

like image 96
David Griffin Avatar answered Sep 30 '22 03:09

David Griffin