如何根據列中的唯一值將 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