如何根据列中的唯一值将dask数据帧拆分为多个分区

Question

我有一个日期列为doc_date的dask数据帧，其范围在12-1-2021到1-2-2022之间。我想重新分区并将这个dask数据帧拆分为26个分区，这样每个分区只有上述日期范围内的1个日期。

以下是我尝试过的：

doc_dates = [dt.strftime("%Y-%m-%d") for dt in pd.date_range('2021-12-08', '2022-01-02')]
predictions_df = predictions_df.set_index('doc_date')
predictions_df = predictions_df.repartition(divisions=sorted(doc_dates))

但我似乎得到了这个错误：

ValueError: left side of old and new divisions are different

Answer 1

问题是，您需要将compute=True传递给dask.dataframe.set_index，以确保数据实际上是按日期排序的，然后才能将日期的排序列表提供给repartition命令：

predictions_df = predictions_df.set_index('doc_date', compute=True)
predictions_df = predictions_df.repartition(divisions=sorted(doc_dates))

或者，您可以将divisions参数用于dask.dataframe.set_index:

predictions_df = predictions_df.set_index(
'doc_date',
divisions=sorted(doc_dates),
compute=True,
)