如何根据列中的唯一值将 dask dataframe 拆分为分区？

Question

我有一个 dask dataframe ，其日期列doc_date在12-1-2021到1-2-2022范围内。 我想将此 dask dataframe 重新分区并拆分为 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

您可以将divisions参数用于dask.dataframe.set_index 。 从 set_index 文档：

部门：列表，可选
用于将新索引拆分为分区的“分隔线”。 对于divisions=[0, 10, 50, 100] ，将有三个 output 分区，其中新索引分别包含 [0, 10)、[10, 50) 和 [50, 100)。 请参阅 https://docs.dask.org/en/latest/dataframe-design.html#partitions。 如果未给出（默认），则通过立即计算数据并查看其值的分布来计算良好的划分。 对于大型数据集，这可能很昂贵。 请注意，如果sorted=True ，则假定指定的分区与数据中的现有分区匹配； 如果这是不正确的，您应该将部门留空并在set_index repartition

设置像您的数据这样的示例：

In [44]: dates = [dt.strftime("%Y-%m-%d") for dt in pd.date_range('2021-12-08', '2022-01-02')]

In [45]: df = ddf.from_pandas(
    ...:     pd.DataFrame({
    ...:         'doc_date': np.random.choice(dates, size=100), 
    ...:         'i': range(100),
    ...:     }),
    ...:     npartitions=10,
    ...: )

您可以对dask.dataframe.set_index使用 divisions 参数：

In [46]: df = df.set_index('doc_date', sorted=False, divisions=list(dates))

您的数据现在将按日期排序，每个日期一个分区：

In [64]: df.partitions[0].compute()
Out[64]:
             i
doc_date
2021-12-08   1
2021-12-08  43
2021-12-08  48
2021-12-08  13

In [65]: df.partitions[1].compute()
Out[65]:
             i
doc_date
2021-12-09  47
2021-12-09  52
2021-12-09  34

In [66]: df.partitions[2].compute()
Out[66]:
             i
doc_date
2021-12-10  32
2021-12-10  10
2021-12-10  23
2021-12-10  63
2021-12-10  66
2021-12-10  75
2021-12-10  93