Как обработать фрейм данных pyspark как значение группы по столбцу

Question

У меня огромный массив данных различных item_id и связанных с ним данных, мне нужно обрабатывать каждую группу с item_id отдельно параллельно, я попытался repartition кадра данных с помощью item_id, используя приведенный ниже код , но кажется, что он все еще обрабатывается в целом, а не кусками

data = sqlContext.read.csv(path='/user/data', header=True)
columns = data.columns    
result = data.repartition('ITEM_ID') \
        .rdd \
        .mapPartitions(lambda iter: pd.DataFrame(list(iter), columns=columns))\
        .mapPartitions(scan_item_best_model)\
        .collect()

также является repartition правильный подход или что-то я делаю не так?

Answer 1

после осмотра я нашел это , которое решает подобную проблему, в конце концов мне пришлось решить ее как

data = sqlContext.read.csv(path='/user/data', header=True)

columns = data.columns

df = data.select("ITEM_ID", F.struct(columns).alias("df"))

df = df.groupBy('ITEM_ID').agg(F.collect_list('df').alias('data'))

df = df.rdd.map(lambda big_df: (big_df['ITEM_ID'], pd.DataFrame.from_records(big_df['data'], columns=columns))).map(
    scan_item_best_model)