如何按列值將 pyspark dataframe 處理為分組

Question

我有一個巨大的 dataframe 不同的item_id及其相關數據，我需要使用item_id並行處理每個組，我嘗試使用下面的整個代碼按item_id repartition dataframe，但似乎它仍然被處理為不是塊

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)