如何根据列中的唯一值将 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,
)