Разделить фрейм данных pyspark на куски и преобразовать в словарь

Question

У меня есть фрейм данных pyspark, который выглядит следующим образом:

+----+--------------------+
|  ID|               Email|
+----+--------------------+
|   1|  sample@example.org|
|   2| sample2@example.org|
|   3|   sampleexample.org|
|   4|   sample@exampleorg|
+----+--------------------+

Что мне нужно сделать, это разделить его на куски, а затем преобразовать эти куски в словари, например:

chunk1
[{'ID': 1, 'Email': 'sample@example.org'}, {'ID': 2, 'Email': 'sample2@example.org'}]

chunk2
[{'ID': 3, 'Email': 'sampleexample.org'}, {'ID': 4, 'Email': 'sample@exampleorg'}]

Я нашел этот пост на SO, но я подумал, что не имеет никакого смысла сначала преобразовывать куски в фрейм данных панд и оттуда в словарь, хотя я мог бы сделать это напрямую. Используя идею из этого поста, я получил следующий код, но не уверен, что это лучший способ сделать это:

columns = spark_df.schema.fieldNames()
chunks = spark_df.repartition(num_chunks).rdd.mapPartitions(lambda iterator: [iterator.to_dict('records')]).toLocalIterator()
for list_of_dicts in chunks:
    # do work locally on list_of_dicts

Answer 1

Вы можете вернуть [[x.asDict() for x in iterator]] в функции mapPartitions (панды не нужны). [x.asDict() for x in iterator] создает список диктов, включающий все строки в одном разделе. затем мы заключаем его в другой список, чтобы он обрабатывался как отдельный элемент с toLocalIterator():

from json import dumps    

num_chunks = 2
chunks = spark_df.repartition(num_chunks).rdd.mapPartitions(lambda iterator: [[x.asDict() for x in iterator]]).toLocalIterator()
for list_of_dicts in chunks:
  print(dumps(list_of_dicts))
#[{"ID": "2", "Email": "sample2@example.org"}, {"ID": "1", "Email": "sample@example.org"}]
#[{"ID": "4", "Email": "sample@exampleorg"}, {"ID": "3", "Email": "sampleexample.org"}]