spark : sum values for same key from a json column while doing a group by based on id in spark dataframe

user_7054529 0
Viewed 81

I have a spark dataframe with id and array of json as columns. I want to do group by based on ID and sum the values for same key in the json column.

Note : I don't want to use explode as i have huge dataset and every json has millions of key and value. If i do explode, my datasize will increase exponentially. I want some solution without explode

My dataframe and schema looks like below :

+-------+----------------------------+
|id     |json                        |
+-------+----------------------------+ 
|a1     |[{key1, 10.0}, {key2, 20.0}]|
|a1     |[{key3, 30.0}, {key1, 40.0}]|
|a2     |[{key1, 10.0}, {key5, 50.0}]|
+-------+----------------------------+

root
 |-- id: string (nullable = true)
 |-- json: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- key: string (nullable = true)
 |    |    |-- value: double (nullable = true)

Expected Output :

+-------+------------------------------------------+
|id     |json                                      |
+-------+------------------------------------------+ 
|a1     |[{key1, 50.0}, {key2, 20.0}, {key3, 30.0}]|
|a2     |[{key1, 10.0}, {key5, 50.0}]              |
+-------+------------------------------------------+
pyspark apache-spark-sql apache-spark
1 Answers

If the array and data are small enough, we can explode the array to create rows for each key-value and aggregate the row. Once aggregated, we can create the array-struct again.

data_sdf.show()

# +---+------------------------+
# |id |json                    |
# +---+------------------------+
# |a1 |[{k1, 10.0}, {k2, 20.0}]|
# |a1 |[{k3, 30.0}, {k1, 40.0}]|
# +---+------------------------+

data2_sdf = data_sdf. \
    withColumn('json_explode', func.explode('json')). \
    select('id', 
           func.col('json_explode.key').alias('key'), 
           func.col('json_explode.value').alias('value')
           ). \
    groupBy('id', 'key'). \
    agg(func.sum('value').alias('value')). \
    withColumn('kv_struct', 
               func.struct(func.col('key').alias('key'), func.col('value').alias('value'))
               )

# +---+---+-----+----------+
# |id |key|value|kv_struct |
# +---+---+-----+----------+
# |a1 |k3 |30.0 |{k3, 30.0}|
# |a1 |k1 |50.0 |{k1, 50.0}|
# |a1 |k2 |20.0 |{k2, 20.0}|
# +---+---+-----+----------+

data2_sdf. \
    groupBy('id'). \
    agg(func.array_sort(func.collect_list('kv_struct')).alias('new_json')). \
    show(truncate=False)

# +---+------------------------------------+
# |id |new_json                            |
# +---+------------------------------------+
# |a1 |[{k1, 50.0}, {k2, 20.0}, {k3, 30.0}]|
# +---+------------------------------------+
user_82795850
Related