如何在大数据帧的每组中有效地随机标记行？

Posted 2023-02-16

【中文标题】如何在大数据帧的每组中有效地随机标记行？

【英文标题】：How to efficiently label randomly rows in each group of a big dataframe?

【发布时间】：2022-01-12 15:40:41

【问题描述】：

我有一个包含 4000 万行的数据框 df。有一个名为group_id 的列来指定行的组标识符。共有2000组。

我想随机标记每个组中的元素，并将此信息添加到 batch 的列 df。例如，如果组 1 包含第 1、2、3、4 和 5 行，那么我选择 (1, 2, 3, 4, 5) 的排列，例如，我们取 (5, 3, 4, 2 , 1)。然后我将值 [5, 3, 4, 2, 1] 分配给这些行的列 batch。

我定义了一个函数func，使用了并行化dummy.Pool，但是速度很慢。您能建议一种更快的方法吗？

import pandas as pd
import numpy as np
import random
import os
from multiprocessing import dummy
import itertools
core = os.cpu_count()
P = dummy.Pool(processes = core)

N = int(4e7)
M = int(2e3) + 1
col_1 = np.random.randint(1, M, N)
col_2 = np.random.uniform(low = 1, high = 5, size = N)
df = pd.DataFrame('group_id': col_1, 'value': col_2)
df.sort_values(by = 'group_id', inplace = True)
df.reset_index(inplace = True, drop = True)

id_ = np.unique(df.group_id)

def func(i):
    idx = df.group_id == i
    m = sum(idx) # count the number of rows in each group
    r = list(range(1, m + 1, 1)) # create an enumeration
    random.shuffle(r) # create a permutation the enumeration
    return(r)
    
order_list = P.map(func, id_)
# merge the list containing permutations
order =  list(itertools.chain.from_iterable(order_list)) 

df['batch'] = order

【问题讨论】：

如果第 2 组包含第 6、7、8、9、10 行 - 它的批次是多少？是再次 1-5 还是 6-10 以某种方式

@JonClements 在您的示例中，组 2 的大小也是 5，因此我们再次创建 (1, 2, 3, 4, 5) 的排列。我的意思是标签在组之间是独立的。这仅取决于组的大小。所有这些都可以让我控制要使用的子数据集的大小。例如，我可以通过idx = df.batch <= 100从4000万行数据集中提取一个子样本。

如果目标只是为每个组创建一个子样本，是否有任何理由不只使用GroupBy sample -> df.groupby('group_id').sample(n=5)？这将随机选择行以确保每组的特定数量和现有行的排列。

你真的想要批次的排列还是只是一个数字？这与 group_id 有何不同？恐怕我不太明白你要做什么实际上会改进其他任何事情，只是增加更多的复杂性？

很高兴听到...我会继续吃完午餐 - 看起来你已经有人解决了你的问题 :)

【参考方案1】：

也许这可以解决你的问题。对组大小进行随机排列。

import numpy as np
import pandas as pd
l = np.repeat([x for x in range(2000)],20000)
df = pd.DataFrame(l, columns=['group'])

df['batch'] = df.groupby('group')['group'].transform(lambda x: np.random.permutation(np.arange(x.size)))