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

Question

我有一个包含40百万行的数据框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)))