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

Question

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

我想随机标记每个组中的元素，并将此信息添加到 df 的列 batch 中。例如，如果第 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

Answer 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)))

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

How to efficiently label randomly rows in each group of a big dataframe?

python

multiprocessing

pandas