如何在使用Dask进行分组聚合后保留分区

Question

在我的应用程序中，我使用groupby对dask dataframe进行聚合，并按特定id排序。

然而，我希望聚合保持分区划分，因为我打算与其他完全分区的dataframe执行连接操作。

import pandas as pd
import numpy as np
import dask.dataframe as dd

df =pd.DataFrame(np.arange(16), columns=['my_data'])
df.index.name = 'my_id'

ddf = dd.from_pandas(df, npartitions=4)
ddf.npartitions
# 4

ddf.divisions
# (0, 4, 8, 12, 15)

aggregated = ddf.groupby('my_id').agg({'my_data': 'count'})
aggregated.divisions
# (None, None)

有没有实现这个目标的方法？

- pygabriel

Answer 1

由于dask需要在分区之间聚合计数，因此您可能无法维护相同的分区。您的数据必须以依赖于数据值的方式移动。

如果您希望确保输出具有多个分区，则可以选择使用agg中的split_out=关键字。

- MRocklin

---

但是索引不会确保例如ID从0到4的数据在同一个分区吗？因此，如果我选择my_id=0的数据组，则可以保证所有元素都在同一个分区中，无需移动任何内容。（实际上，我可以使用apply方法来获得完全相同分区的结果）。我想知道是否可以使用agg方法做到同样的效果。 - pygabriel

啊，我明白了，我没有注意到你的数据框是以 my_id 为索引的。 - MRocklin