Question

我有一个Dask DataFrames，其中包含不唯一的索引（client_id）。重新分区和重置索引最终会导致分区非常不均匀-有些分区只包含几行，数千行。例如以下代码：

for p in range(ddd.npartitions):
    print(len(ddd.get_partition(p)))

打印出类似这样的内容：

我的DataFrame是一键编码的，具有500多个列。较大的分区不适合内存。我想对DataFrame进行重新分区以使其具有甚至大小的分区。您知道一种有效的方法吗？

编辑1

简单复制：

df = pd.DataFrame({'x':np.arange(0,10000),'y':np.arange(0,10000)})
df2 = pd.DataFrame({'x':np.append(np.arange(0,4995),np.arange(5000,10000,1000)),'y2':np.arange(0,10000,2)})
dd_df = dd.from_pandas(df, npartitions=10).set_index('x')
dd_df2= dd.from_pandas(df2, npartitions=5).set_index('x')
new_ddf=dd_df.merge(dd_df2, how='right')
#new_ddf = new_ddf.reset_index().set_index('x')
#new_ddf = new_ddf.repartition(npartitions=2)
new_ddf.divisions
for p in range(new_ddf.npartitions):
    print(len(new_ddf.get_partition(p)))

注意最后一个分区（一个元素）：

即使我们取消注释行的注释，分区的大小仍然不均匀。

编辑II：Walkoround

简单的wlakoround可以通过以下代码实现。有没有更精致的方法可以做到这一点（更多的方式是Dask）？

def repartition(ddf, npartitions=None):
    MAX_PART_SIZE = 100*1024

    if npartitions is None:
        npartitions = ddf.npartitions

    one_row_size = sum([dt.itemsize for dt in ddf.dtypes])
    length = len(ddf)

    requested_part_size = length/npartitions*one_row_size
    if requested_part_size <= MAX_PART_SIZE:
        np = npartitions
    else:
        np = length*one_row_size/MAX_PART_SIZE

    chunksize = int(length/np)


    vc = ddf.index.value_counts().to_frame(name='count').compute().sort_index()

    vsum = 0
    divisions = [ddf.divisions[0]]
    for i,v in vc.iterrows():
        vsum+=v['count']
        if vsum > chunksize:
            divisions.append(i)
            vsum = 0
    divisions.append(ddf.divisions[-1])


    return ddf.repartition(divisions=divisions, force=True)

Answer 1

您是正确的，.repartition不会解决问题，因为它不处理任何用于计算除法的逻辑，而只是尝试尽可能地合并现有分区。这是我针对相同问题提出的解决方案：

def _rebalance_ddf(ddf):
    """Repartition dask dataframe to ensure that partitions are roughly equal size.

    Assumes `ddf.index` is already sorted.
    """
    if not ddf.known_divisions:  # e.g. for read_parquet(..., infer_divisions=False)
        ddf = ddf.reset_index().set_index(ddf.index.name, sorted=True)
    index_counts = ddf.map_partitions(lambda _df: _df.index.value_counts().sort_index()).compute()
    index = np.repeat(index_counts.index, index_counts.values)
    divisions, _ = dd.io.io.sorted_division_locations(index, npartitions=ddf.npartitions)
    return ddf.repartition(divisions=divisions)

内部函数sorted_division_locations可以满足您的需求，但仅适用于类似于列表的实际列表，而不适用于惰性dask.dataframe.Index。这样可以避免在有很多重复项的情况下提取完整索引，而只是获取计数并从中进行本地重构。

如果数据框太大，甚至索引也无法容纳在内存中，那么您就需要做一些更聪明的事情。

重新分区Dask DataFrame以获得均匀分区

1 个答案: