以大约相等的计算成本将不等大小的 pandas/numpy 数组分箱

Question

我遇到了一个问题，即必须跨多个内核处理数据。让 df 成为 Pandas DataFrameGroupBy (size()) 对象。每个值代表每个 GroupBy 对核心的计算“成本”。如何将 df 划分为 大小不等 且计算成本相同 （近似）的 n 箱？

import pandas as pd
import numpy as np
size = 50
rng = np.random.default_rng(2021)
df = pd.DataFrame({
    "one": np.linspace(0, 10, size, dtype=np.uint8),
    "two": np.linspace(0, 5, size, dtype=np.uint8),
    "data": rng.integers(0, 100, size)
})
groups = df.groupby(["one", "two"]).sum()

df
    one  two  data
0     0    0    75
1     0    0    75
2     0    0    49
3     0    0    94
4     0    0    66
...
45    9    4    12
46    9    4    97
47    9    4    12
48    9    4    32
49   10    5    45

人们通常将数据集拆分为 n 个箱，例如下面的代码。但是，将数据集分成 n 等份是不可取的，因为核心接收非常不平衡的工作负载，例如205 与 788。

n = 4
bins = np.array_split(groups, n) # undesired

[b.sum() for b in bins]  #undesired
[data    788
dtype: int64, data    558
dtype: int64, data    768
dtype: int64, data    205
dtype: int64]

理想的解决方案是将数据拆分为大小不等且总和值大致相等的 bin。 IE。 abs(743-548) = 195 之间的差异小于之前的方法abs(205-788) = 583。差异应尽可能小。一个简单的列表示例，说明它应该如何实现：

# only an example to demonstrate desired functionality
example = [[[10, 5], 45], [[2, 1], 187], [[3, 1], 249], [[6, 3], 262]], [[[9, 4], 153], [[4, 2], 248], [[1, 0], 264]], [[[8, 4], 245], [[7, 3], 326]], [[[5, 2], 189], [[0, 0], 359]]

[sum([size for (group, size) in test]) for test in t]  # [743, 665, 571, 548]

在 pandas 或 numpy 中是否有更有效的方法将数据集拆分为 bin？

拆分/装箱 GroupBy 对象很重要，以与np.array_split() 返回的类似方式访问数据。

Answer 1

我认为已经找到了一个好的方法。向同事致谢。

这个想法是对组大小进行排序（按降序排列），并以“向后 S”模式将组放入箱中。让我用一个例子来说明。假设n = 3（箱数）和以下数据：

groups
    data
0    359
1    326
2    264
3    262
4    249
5    248
6    245
7    189
8    187
9    153
10    45

这个想法是把一组放在一个垃圾箱中，在垃圾箱之间以“向后 S”模式“从右到左”（反之亦然）。 bin 0 中的第一个元素，bin 1 中的第二个元素等。然后在到达最后一个 bin 时向后退：bin 2 中的第四个元素，bin 1 中的第五个元素等。请参阅下面如何将元素按组号放入 bin在括号中。这些值是组大小。

 Bins:  |    0    |    1    |    2    |
        |  359 (0)|  326 (1)|  264 (2)|  
        |  248 (5)|  249 (4)|  262 (3)|
        |  245 (6)|  189 (7)|  187 (8)|
        |         |   45(10)|  153 (9)|

这些 bin 将具有大致相同数量的值，因此，计算“成本”也大致相同。垃圾箱大小为：[852, 809, 866] 任何感兴趣的人。我已经尝试过真实世界的数据集，并且这些垃圾箱的大小相似。不能保证所有数据集的 bin 大小相似。

代码可以变得更高效，但这足以让想法变得清晰：

n = 3
size = 50
rng = np.random.default_rng(2021)
df = pd.DataFrame({
    "one": np.linspace(0, 10, size, dtype=np.uint8),
    "two": np.linspace(0, 5, size, dtype=np.uint8),
    "data": rng.integers(0, 100, size)
})

groups = df.groupby(["one", "two"]).sum()
groups = groups.sort_values("data", ascending=False).reset_index(drop=True)

bins = [[] for i in range(n)]
backward = False
i = 0
for group in groups.iterrows():
    bins[i].append(group)
    i = i + 1 if not backward else i - 1
    if i == n:
        backward = True
        i -= 1
    if i == -1 and backward:
        backward = False
        i += 1


[sum([size[0] for (group, size) in bin]) for bin in bins]