Create a complex dataframe on (py)spark - pyspark

Question

I want to know how I can create the following “complex” json structure in (py)Spark (2.3.2):
Test data set:
df = sc.parallelize([
[1, 'a', 'x1'],
[1, 'a', 'x2'],
[2, 'b', 'y1'],
[2, 'b', 'y2'],
[3, 'c', 'z'],
]).toDF('id: integer, field1: string, field2: string').cache()
Code:
import pyspark.sql.functions as F
out = df.groupBy('id').agg(F.to_json(F.create_map(
F.lit('key1'),
F.col('field1'),
F.lit('info'),
F.collect_list(F.create_map(
F.lit('number'),
F.col('id'),
F.lit('key2'),
F.col('field2')
))
)).alias('info'))
What my target json structure is, is that I have a data set like this:
[
(id) 1, (info) {“key1”: ‘a’, “info”: [{“number”: 1, “key2”: “x1”}, {“number”: 1, “key2”: “x1”}],
(id) 2, (info) {“key1”: ‘b’, “info”: [{“number”: 2, “key2”: “y1”}, {“number”: 1, “key2”: “x2”}],
(id) 3, (info) {“key1”: ‘c’, “info”: [{“number”: 3, “key2”: “z”}]
]
How could I achieve this? (Can I achieve this?)
As I'm always getting the following error:
org.apache.spark.sql.AnalysisException:
cannot resolve 'map('key1', `field1`, 'info', collect_list(map('number',
CAST(`id` AS STRING), 'key2', CAST(`field2` AS STRING))))'
due to data type mismatch: The given values of function map should all be the same type,
but they are [string, array<map<string,string>>]
What I understand from this error is that field1 is a string, and the value of 'info' is not. But that's the way I want it to be...
So, could I achieve this another way?
Thanks!

Answer 1

I found one (hackish) way to do things... I don't like it very much but seeing that no-one in this community posted an answer, I am starting to think it's not that easy.
So first of all, I split the "big" aggregation in 2:
out = df.groupBy('id', 'field1').agg(F.to_json(F.create_map(
F.lit('key1'),
F.col('field1'),
F.lit('info'),
F.lit('%%replace%%')
)).alias('first'), F.to_json( F.collect_list(F.create_map(
F.lit('number'),
F.col('id'),
F.lit('key2'),
F.col('field2')
))
).alias('second'))
This will generate the following table:
+---+------+---------------------------------+-------------------------------------------------------+
|id |field1|first |second |
+---+------+---------------------------------+-------------------------------------------------------+
|3 |c |{"key1":"c","info":"%%replace%%"}|[{"number":"3","key2":"z"}] |
|2 |b |{"key1":"b","info":"%%replace%%"}|[{"number":"2","key2":"y1"},{"number":"2","key2":"y2"}]|
|1 |a |{"key1":"a","info":"%%replace%%"}|[{"number":"1","key2":"x1"},{"number":"1","key2":"x2"}]|
+---+------+---------------------------------+-------------------------------------------------------+
And now you combine them together:
df2 = out.withColumn('final', F.expr('REPLACE(first, \'"%%replace%%"\', second)')).drop('first').drop('second')
df2.show(10, False)
+---+------+---------------------------------------------------------------------------+
|id |field1|final |
+---+------+---------------------------------------------------------------------------+
|3 |c |{"key1":"c","info":[{"number":"3","key2":"z"}]} |
|2 |b |{"key1":"b","info":[{"number":"2","key2":"y1"},{"number":"2","key2":"y2"}]}|
|1 |a |{"key1":"a","info":[{"number":"1","key2":"x1"},{"number":"1","key2":"x2"}]}|
+---+------+---------------------------------------------------------------------------+
A bit unorthodox, but no complaints from Spark :)