How to aggregate 2 columns into map in pyspark

user_11803907 0
Viewed 586

I have a DataFrame like this

a = spark.createDataFrame([['Alice', '2020-03-03', '1'], ['Bob', '2020-03-03', '1'], ['Bob', '2020-03-05', '2']], ['name', 'dt', 'hits'])
a.show()
+-----+----------+----+
| name|        dt|hits|
+-----+----------+----+
|Alice|2020-03-03|   1|
|  Bob|2020-03-03|   1|
|  Bob|2020-03-05|   2|
+-----+----------+----+

I want to aggregate dt and hits Columns into map -

+-----+------------------------------------+
| name|    map                             |
+-----+------------------------------------+
|Alice|   {'2020-03-03': 1, '2020-03-05':2}|
|  Bob|   {'2020-03-03': 1}                |
+-----+------------------------------------+

But this code throw an exception:

from pyspark.sql import functions as F
a = a.groupBy(F.col('name')).agg(F.create_map(F.col('dt'), F.col('hits')))

Py4JJavaError: An error occurred while calling o2920.agg.
: org.apache.spark.sql.AnalysisException: expression '`dt`' is neither present in the group by, nor is it an aggregate function. Add to group by or wrap in first() (or first_value) if you don't care which value you get.;;
Aggregate [name#1329], [name#1329, map(dt#1330, hits#1331) AS map(dt, hits)#1361]
+- LogicalRDD [name#1329, dt#1330, hits#1331], false

What am i doing wrong?

pyspark python
1 Answers

For spark2.4+, you can use map_from_arrays like this: 

from pyspark.sql import functions as F

a.groupBy("name").agg(F.map_from_arrays(F.collect_list("dt"),\
                     F.collect_list("hits")).alias("map")).show(truncate=False)

#+-----+----------------------------------+
#|name |map                               |
#+-----+----------------------------------+
#|Bob  |[2020-03-03 -> 1, 2020-03-05 -> 2]|
#|Alice|[2020-03-03 -> 1]                 |
#+-----+----------------------------------+
user_101073890
Related