将函数应用于Dask中的分组数据框:如何在函数中将分组的Dataframe指定为参数？

Question

我有一个dask dataframe按索引(first_name)分组.

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将每行与数据帧中的所有行进行比较,并将结果保存在每行的列表中).

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

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.注意,DF在contraster函数中必须是分组的数据帧.

我想使用以下内容:

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

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

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

Answer 2

稍微猜测一下，我认为以下是您所追求的。

def mapper(d):

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

df.groupby('first_name').apply(mapper).compute()

应用于您的数据，您将获得：

   ID first_name  last_name   out
2   X      Danae      Smith   [X]
4  12      Jacke       Toro  [12]
0   X       Jake   Del Toro   [X]
1   U       John     Foster   [U]
5  13        Jon    Froster  [13]
3   Y    Beatriz  Patterson   [Y]

即，因为您按first_name分组，所以每个组只包含一个仅与自身匹配的项目。

但是，如果您有一些位于多行中的first_name值，您将获得匹配项：

entities = pd.DataFrame(
    {'first_name':['Jake','Jake', 'Jake', 'John'],
     'last_name': ['Del Toro', 'Toro', 'Smith'
                   'Froster'],
     'ID':['Z','U','X','Y']})

输出：

  ID first_name last_name     out
0  Z       Jake  Del Toro  [Z, U]
1  U       Jake      Toro  [Z, U]
2  X       Jake     Smith     [X]
3  Y       John   Froster     [Y]

如果您不需要first_name上的完全匹配，那么您可能需要按 first_name 排序/设置索引并以类似的方式使用。在这种情况下，您将需要修改您的问题。map_partitions