如何使用reduceByKey（pyspark）拥有嵌套结构？

Question

I'm working with spark (pyspark) on a data set which I want to partition based on 3 values and write back to S3. Data set looks like below -

customerId, productId, createDate

I would like to partition this data by customerId then productId then createDate . So when I write this partitioned data to s3, it should have the below structure -

customerId=1
  productId='A1'
    createDate=2019-10
    createDate=2019-11
    createDate=2019-12
  productId='A2'
    createDate=2019-10
    createDate=2019-11
    createDate=2019-12

below is the code that I'm using to create the partition.

rdd = sc.textFile("data.json")  #sc is spark context
r1.map(lambda r: (r["customerId"], r["productId"],r["createDate"])).distinct().map(lambda r: (r[0], ([r[1]],[r[2]]))).reduceByKey(lambda a, b: (a[0] + b[0],a[1] + b[1])).collect()

[('1', ([A1,A2], ['2019-12', '2019-11', '2019-10', '2019-12', '2019-11', '2019-10']))]

This code does gives me a flat structure and not the nested which I mentioned. Is it possible to transform the way I describe. any pointer is highly appretiated.

Answer 1

first read your JSON file to dataframe.

import json
a=[json.dumps("/data.json")]
jsonRDD = sc.parallelize(a)
df = spark.read.json(jsonRDD)

then use groupby and collectlist to get the desired format.

import pyspark.sql.functions as func
df.groupby('customerId','productId').agg(func.collectList('createDate')).collect()