Scala spark how to interact with a List[Option[Map[String, DataFrame]]]

user_14567747 0
Viewed 291

I'm trying to interact with this List[Option[Map[String, DataFrame]]] but I'm having a bit of trouble.

Inside it has something like this:

customer1 -> dataframeX 
customer2 -> dataframeY 
customer3 -> dataframeZ

Where the customer is an identifier that will become a new column.

I need to do an union of dataframeX, dataframeY and dataframeZ (all df have the same columns). Before I had this:

map(_.get).reduce(_ union _).select(columns:_*)

And it was working fine because I only had a List[Option[DataFrame]] and didn't need the identifier but I'm having trouble with the new list. My idea is to modify my old mapping, I know I can do stuff like "(0).get" and that would bring me "Map(customer1 -> dataframeX)" but I'm not quite sure how to do that iteration in the mapping and get the final dataframe that is the union of all three plus the identifier. My idea:

map(/*get identifier here along with dataframe*/).reduce(_ union _).select(identifier +: columns:_*)

The final result would be something like:

-------------------------------
|identifier | product  |State | 
-------------------------------
|  customer1|  prod1   |  VA  |
|  customer1|  prod132 |  VA  |
|  customer2|  prod32  |  CA  | 
|  customer2|  prod51  |  CA  |
|  customer2|  prod21  |  AL  |
|  customer2|  prod52  |  AL  |
-------------------------------
apache-spark mapping reduce list scala
2 Answers

You could use collect to unnest Option[Map[String, Dataframe]] to Map[String, DataFrame]. To put an identifier into the column you should use withColumn. So your code could look like:

import org.apache.spark.sql.functions.lit

val result: DataFrame = frames.collect {
    case Some(m) =>
      m.map {
        case (identifier, dataframe) => dataframe.withColumn("identifier", lit(identifier))
      }.reduce(_ union _)
  }.reduce(_ union _)
user_13033240

Something like this perhaps?

list
  .flatten 
  .flatMap { 
    _.map { case (id, df) => 
      df.withColumn("identifier", id) } 
  }.reduce(_ union _)
user_42545170
Related