Custom aggregation to a JSON in pyspark

user_4447106 0
Viewed 114

This is my current table:

User Id Product Amount 1 Amount 2 Amount 3
1 A 100 200 300
1 B 200 300 400
2 A 500 600 700

This is the output I'm looking for:

User Id Amount 1 Amount 2 Amount 3
1 {"A": 100, "B": 200} {"A": 200, "B": 300} {"A": 300, "B": 400}
2 {"A": 500} {"A": 600} {"A": 700}

I understand that I should be using a User Defined Aggregate Function, but am unable to understand how to implement them via PySpark.

Any help would be greatly appreciated.

pyspark apache-spark-sql apache-spark python
1 Answers

You can use to_json after grouping by User Id and creating map column:

from pyspark.sql import functions as F

df1 = df.groupBy("User Id").agg(*[
    F.to_json(
        F.map_from_entries(F.collect_list(F.struct(F.col("Product"), F.col(c))))
    ).alias(c)
    for c in df.columns[2:]
])

df1.show()
#+-------+-----------------+-----------------+-----------------+
#|User Id|         Amount 1|         Amount 2|         Amount 3|
#+-------+-----------------+-----------------+-----------------+
#|      1|{"A":100,"B":200}|{"A":200,"B":300}|{"A":300,"B":400}|
#|      2|        {"A":500}|        {"A":600}|        {"A":700}|
#+-------+-----------------+-----------------+-----------------+
user_13865510
Related