dask groupby 不合并分区

我有一组数据，我想要对其进行一些简单的 groupby/count 操作，但我似乎无法使用 dask 来完成此操作。

我很可能不理解 dask 中执行 groupby/reduce 的方式，特别是当索引位于分组键中时。所以我将用玩具数据来说明我的问题。

首先我创建一个包含 3 列的数据框。

import pandas as pd
import numpy as np

np.random.seed(0)
df = pd.DataFrame(
    {"A": np.random.randint(6, size=20),
     "B": np.random.randint(6, size=20),
     "C": np.random.rand(20)}
 )
 df = df.set_index("A")

所以我有一个带有索引和 2 列的数据框。在熊猫中我会这样做：

result = df.groupby(["A", "B"]).C.count().reset_index().set_index("A")

最后我想将结果保存在镶木地板文件中。

现在让我们进入 dask，我可以执行基本相同的操作：

import dask.dataframe as dd
ddf = dd.from_pandas(df, npartitions=4)
result = ddf.groupby(["A", "B"]).C.count().reset_index().set_index("A")

调用compute会得到相同的结果。然而，当检查图表时我发现​​：

GroupBy/Count 的计算图 https://i.stack.imgur.com/ObKPU.png

我们看到一切都被简化为一个分区。我可以理解，当要分组的数据分布在不同的分区中或者实际分组键的数量很少时，这或多或少是必要的。

但就我而言，索引位于分区中，因此我希望得到一个完全并行的图。此外，分组键的数量与初始行的数量具有相同的数量级（除以因子2或3）。

就像这段代码得到的一样：

result = ddf.map_partitions(
     lambda x: x.groupby(
         [x.index, x.B]
     ).C.count().reset_index().set_index("A")
)

给出了下图：并行计算图 https://i.stack.imgur.com/fnlFg.png

那么有没有办法用正常的 groupby/reduce dask 函数获得这个并行图？

当分组键的数量非常大时，这一点非常重要。