How do I merge rows of maps into one map in Spark?

Viewed 23

I currently have a dataframe:

df
.groupBy($"letters")
.agg(collect_list($"numbers").as("numbers"))
.select(map($"letters",$"numbers").as("data"))
.agg(collect_list($"data").as("data"))
.select(to_json($"data").as("output"))
.show(false)

+------------------------------------------+
|output                                    |
+------------------------------------------+
|[{"abc":["123","456"]},{"def":["123"]}]   |
+------------------------------------------+

How can I get it into this format?

+------------------------------------------+
|output                                    |
+------------------------------------------+
|{"abc":["123","456"],"def":["123"]}       |
+------------------------------------------+

So that basically it is one map, with no [] brackets at the ends

In other words, it is currently

res34: org.apache.spark.sql.DataFrame = [data: array<map<string,array<string>>>]
1 Answers

Assuming your dataset is this:

+-------------------------------------+
|output                               |
+-------------------------------------+
|[{abc -> [123, 456]}, {def -> [123]}]|
+-------------------------------------+

That has been created through:

var ds = spark.sparkContext.parallelize(Seq(
  List(Map("abc" -> Array("123", "456")), Map("def" -> Array("123"))),
)).toDF("output")

We get:

+-------------------------------------+----------+-------------------+---------------------------------+
|output                               |keys      |values             |map                              |
+-------------------------------------+----------+-------------------+---------------------------------+
|[{abc -> [123, 456]}, {def -> [123]}]|[abc, def]|[[123, 456], [123]]|{abc -> [123, 456], def -> [123]}|
+-------------------------------------+----------+-------------------+---------------------------------+

Through:

ds = ds
  .withColumn("keys", expr("transform(output, x -> map_keys(x)[0])"))
  .withColumn("values", expr("transform(output, x -> flatten(map_values(x)))"))
  .withColumn("map", map_from_arrays(col("keys"), col("values")))

I could not think of a better solution, the code is also self explanatory so I don't think there is a need for comments, hope it works for you, good luck!

Related