Create DataFrame grouping columns in map type

user_3624549 0
Viewed 7677

My DataFrame has the following structure:

df = spark.createDataFrame(
    [('B', 'a', 10),
     ('B', 'b', 20),
     ('C', 'c', 30)],
    ['Brand', 'Type', 'Amount'])
df.show()
# +-----+----+------+
# |Brand|Type|Amount|
# +-----+----+------+
# |    B|   a|    10|
# |    B|   b|    20|
# |    C|   c|    30|
# +-----+----+------+

I want to reduce the amount of rows by grouping type and amount into one single column of type: Map. So Brand will be unique and MAP_type_AMOUNT will have key,value for each type amount combination.

I think, spark.sql might have some functions to do it, or do I have to use RDD and make my "own" conversion to map type?

Expected output:

---------------------------
| Brand | MAP_type_AMOUNT |
---------------------------
|  B    |  {a: 10, b:20}  |
|  C    |  {c: 30}        |
---------------------------
pyspark apache-spark-sql apache-spark dictionary python
4 Answers

Using collect_list and map_from_arrays together can achieve this

import pyspark.sql.functions as F

df_converted = (
    df.groupBy('Brand')
    .agg(
        F.collect_list('type').alias('type'),
        F.collect_list('amount').alias('amount'),
    )
    .withColumn('MAP_type_AMOUNT', F.map_from_arrays('type', 'amount'))
    .drop('type', 'amount')
)

Output

+-----+------------------+
|Brand|   MAP_type_AMOUNT|
+-----+------------------+
|    C|         [c -> 30]|
|    B|[b -> 20, a -> 10]|
+-----+------------------+
user_71592290

This is a working approach to create map from separate columns:

import pyspark.sql.functions as F

df_converted = df.groupBy('Brand').agg(
    F.map_from_entries(F.collect_set(F.struct('Type', 'Amount'))).alias('MAP_type_AMOUNT')
)

df_converted.show()
# +-----+------------------+
# |Brand|   MAP_type_AMOUNT|
# +-----+------------------+
# |    B|{b -> 20, a -> 10}|
# |    C|         {c -> 30}|
# +-----+------------------+
user_27535010
Related