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

Question

我有一个 dask 数据框，其日期列 doc_date 在 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,
)

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

How to split dask dataframe into partitions based on unique values in a column?

python

dataframe

dask

dask-distributed

dask-dataframe