Фильтрация данных массива json в фрейме искры

Question

У меня есть искровой фрейм данных, который я конвертирую в формат JSON:

json = df.toJSON().collect()
print(json)

['{"lot_number":"4f19-9deb-0ef861c1a6a1","recipients":[{"account":"45678765457876545678","code":"user1","status":"pending"},{"account":"12354567897545678","code":"error2","status":"pending"}]}', 

'{"lot_number":"09ad-451e-8fb1-50bc185ef02f","recipients":[{"account":"4567654567876545678","code":"user3","status":"pending"},{"account":"12354567876545678","code":"user2","status":"pending"}]}']

Мне нужно отфильтровать данные из массива, то есть всех получателей, чей код "user1".

Я ожидаю этого результата:

['{"lot_number":"4f19-9deb-0ef861c1a6a1","recipients":[{"account":"45678765457876545678","code":"user1","status":"pending"}' 
]

Может ли кто-нибудь помочь отфильтровать данные, как показано выше?

Answer 1

Сначала вам нужно будет преобразовать string в list в dict объекты.

import json

json_rdd = df.toJSON().collect()
json_ls = [json.loads(x) for x in json_rdd]

# Now you can filter using "user1"
final_json_ls = [x for x in json_ls if x.get("recipients")[0].get("code") == "user1"]

Если у вас несколько получателей:

new_list = list()
for lot in json_ls:
    recs = lot.get('recipients')
    lot_recipients = [rec for rec in recs if rec.get("code") == "user1"]
    if lot_recipients:
        new_list.append({"lot_number": lot.get('lot_number'),
                         "recipients": lot_recipients})

# OUTPUT
# [{'lot_number': u'4f19-9deb-0ef861c1a6a1', 'recipients': [{u'status': u'pending', u'account': u'45678765457876545678', u'code': u'user1'}]}]

А поскольку вы хотите преобразовать его обратно в json для отправки POST-запросов:

for ls in new_list:
    lot = ls.get("lot_number")
    url = "test.com/api/v1/notify/request/"+ batch
    response = requests.put(url, data=item, headers=headers) 
    print(response.text)