DASK DataFrame：`set_index`可以将单个索引放入多个分区吗？

问题描述

根据经验，无论何时在DaskDataFrame上set_index，Dask都会将具有相等索引的行放入单个分区中，即使这会导致严重的分区不平衡。

这里有一个演示：

import pandas as pd
import dask.dataframe as dd

users = [1]*1000 + [2]*1000 + [3]*1000

df = pd.DataFrame({'user': users})
ddf = dd.from_pandas(df, npartitions=1000)

ddf = ddf.set_index('user')

counts = ddf.map_partitions(lambda x: len(x)).compute()
counts.loc[counts > 0]
# 500    1000
# 999    2000
# dtype: int64

然而，我发现任何地方都不能保证这种行为。

我曾试图自己筛选代码，但放弃了。我相信这些相互关联的函数中的一个可能包含答案：

• set_index
• set_partitions
• rearrange_by_column
• rearrange_by_column_tasks
• SimpleShuffleLayer

当set_index时，单个索引是否永远不能位于两个不同的分区中？如果不是，则此属性在什么条件下有效？

---

赏金：我将奖励来自可靠消息来源的答案。例如，引用实现以显示此属性必须保持。

---

解决方案

单个索引永远不能位于两个不同的分区中吗？

不，这当然是允许的。达斯克甚至会打算让这件事发生。但是，由于set_index中的bug，所有数据仍将保存在一个分区中。

一个极端的例子(除了一行之外，每一行都是相同的值)：

In [1]: import dask.dataframe as dd
In [2]: import pandas as pd
In [3]: df = pd.DataFrame({"A": [0] + [1] * 20})
In [4]: ddf = dd.from_pandas(df, npartitions=10)
In [5]: s = ddf.set_index("A")
In [6]: s.divisions
Out[6]: (0, 0, 0, 0, 0, 0, 0, 1)

如您所见，DASK打算在多个分区之间拆分0。然而，当洗牌实际发生时，所有0仍在一个分区中结束：

In [7]: import dask
In [8]: dask.compute(s.to_delayed())  # easy way to see the partitions separately
Out[8]: 
([Empty DataFrame
  Columns: []
  Index: [],
  Empty DataFrame
  Columns: []
  Index: [],
  Empty DataFrame
  Columns: []
  Index: [],
  Empty DataFrame
  Columns: []
  Index: [],
  Empty DataFrame
  Columns: []
  Index: [],
  Empty DataFrame
  Columns: []
  Index: [],
  Empty DataFrame
  Columns: []
  Index: [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]],)

这是因为code决定一行属于哪个输出分区不考虑divisions中的重复项。它将divisions视为一个序列，使用searchsorted和side="right"，因此所有数据始终位于最后一个分区中。

问题解决后，我将更新此答案。