Spark: написать JSON несколько файлов из DataFrame на основе разделения по значению столбца

Question

Предположим, что у меня есть этот DataFrame (df):

user food  affinity 
'u1' 'pizza'  5 
'u1' 'broccoli' 3 
'u1' 'ice cream' 4 
'u2' 'pizza'  1 
'u2' 'broccoli' 3 
'u2' 'ice cream' 1 

То есть каждый пользователь имеет определенную (вычисленный) сродством к ряду продуктов. DataFrame построен из нескольких, что мне нужно сделать, это создать файл JSON для каждого пользователя с их сходством. Например, для пользователя «u1», я хочу иметь файл для пользователя «u1», содержащего

[ 
    {'food': 'pizza', 'affinity': 5}, 
    {'food': 'broccoli', 'affinity': 3}, 
    {'food': 'ice cream', 'affinity': 4}, 
] 

Это повлечет за собой разделение DataFrame по пользователю, и я не могу придумать способ сделать это, как написание файла JSON будет достигнута, для полного DataFrame, с

df.write.json(<path_to_file>) 

Answer 1

Вы можете partitionBy (это даст вам один каталог и, возможно, несколько файлов для каждого пользователя):

df.write.partitionBy("user").json(<path_to_file>) 

или repartition и partitionBy (это даст вам один каталог и один файл для каждого пользователя):

df.repartition(col("user")).write.partitionBy("user").json(<path_to_file>) 

К сожалению, ни один из вышеперечисленных даст вам массив JSON.

Если вы используете Спарк 2.0 вы можете попробовать с сборным списке первым:

df.groupBy(col("user")).agg(
    collect_list(struct(col("food"), col("affinity"))).alias("affinities") 
) 

и partitionBy на записи, как и прежде.

До 2.0 вам нужно будет использовать API RDD, но это зависит от языка.