python - 将函数应用于 Dask : How do you specify the grouped Dataframe as argument in the function? 中的分组数据帧

Question

我有一个按索引 (first_name) 分组的 dask 数据帧。

import pandas as pd
import numpy as np

from multiprocessing import cpu_count

from dask import dataframe as dd
from dask.multiprocessing import get 
from dask.distributed import Client


NCORES = cpu_count()
client = Client()

entities = pd.DataFrame({'first_name':['Jake','John','Danae','Beatriz', 'Jacke', 'Jon'],'last_name': ['Del Toro', 'Foster', 'Smith', 'Patterson', 'Toro', 'Froster'], 'ID':['X','U','X','Y', '12','13']})

df = dd.from_pandas(entities, npartitions=NCORES)
df = client.persist(df.set_index('first_name'))

(显然现实生活中的entities是几千行)

我想将用户定义的函数应用于每个分组的数据帧。我想将每一行与组中的所有其他行进行比较(类似于 Pandas compare each row with all rows in data frame and save results in list for each row )。

以下是我尝试应用的功能:

def contraster(x, DF):
    matches = DF.apply(lambda row: fuzz.partial_ratio(row['last_name'], x) >= 50, axis = 1) 
    return [i for i, x in enumerate(matches) if x]

对于测试 entities 数据框，您可以照常应用该函数:

entities.apply(lambda row: contraster(row['last_name'], entities), axis =1)

而预期的结果是:

Out[35]: 
0    [0, 4]
1    [1, 5]
2       [2]
3       [3]
4    [0, 4]
5    [1, 5]
dtype: object

当entities很大时，解决方案是使用dask。注意 contraster 函数中的 DF 必须是分组数据帧。

我正在尝试使用以下内容:

df.groupby('first_name').apply(func=contraster, args=????)

但是我应该如何指定分组数据框(即 contraster 中的 DF？)

Answer 1

您提供给 groupby-apply 的函数应将 Pandas 数据帧或系列作为输入，理想情况下返回一个(或标量值)作为输出。额外的参数很好，但它们应该是次要的，而不是第一个参数。这在 Pandas 和 Dask 数据帧中都是一样的。

def func(df, x=None):
    # do whatever you want here
    # the input to this function will have all the same first name
    return pd.DataFrame({'x': [x] * len(df),
                         'count': len(df),
                         'first_name': df.first_name})

然后您可以正常调用 df.groupby

import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame({'first_name':['Alice', 'Alice', 'Bob'],
                   'last_name': ['Adams', 'Jones', 'Smith']})

ddf = dd.from_pandas(df, npartitions=2)

ddf.groupby('first_name').apply(func, x=3).compute()

这将在 pandas 或 dask.dataframe 中产生相同的输出

   count first_name  x
0      2      Alice  3
1      2      Alice  3
2      1        Bob  3