python – Dask：使用groupby获取具有最大值的行

2024年1月24日 251次阅读

使用变换可以在Pandas中解决相同的问题,如
here所述

使用dask唯一正在工作的
solution我发现使用merge.我想知道是否还有其他方法可以实现它. 最佳答案 首先,我想在原始问题中重写引用的脚本,以确保我已理解其意图.据我所知,正如我在下面的重写所示,你基本上想要一种方法来提取具有最高计数cnt值的值,用于foo和bar的每个唯一配对.下面是大致如何使用Pandas完成引用的脚本.

# create an example dataframe
df = pd.DataFrame({
        'foo' : ['MM1', 'MM1', 'MM1', 'MM2', 'MM2', 'MM2', 'MM4', 'MM4', 'MM4'],
        'bar' : ['S1', 'S1', 'S3', 'S3', 'S4', 'S4', 'S2', 'S2', 'S2'],
        'cnt' : [3, 2, 5, 8, 10, 1, 2, 2, 7],
        'val' : ['a', 'n', 'cb', 'mk', 'bg', 'dgb', 'rd', 'cb', 'uyi'],
    })


grouped_df = (df.groupby(['foo', 'bar'])            # creates a double nested indices
                .agg({'cnt': 'max'})                # returns max value from each grouping
                .rename(columns={'cnt': 'cnt_max'}) # renames the col to avoid conflicts on merge later
                .reset_index())                     # makes the double nested indices columns instead

merged_df = pd.merge(df, grouped_df, how='left', on=['foo', 'bar'])

# note: I believe a shortcoming here is that if ther eis more than one match, this would 
# return multiple results for some pairings...
final_df = merged_df[merged_df['cnt'] == merged_df['cnt_max']]

现在,下面是我对Dask-ready版本的看法.查看评论以进行详细说明.

# create an example dataframe
df = pd.DataFrame({
        'foo' : ['MM1', 'MM1', 'MM1', 'MM2', 'MM2', 'MM2', 'MM4', 'MM4', 'MM4'],
        'bar' : ['S1', 'S1', 'S3', 'S3', 'S4', 'S4', 'S2', 'S2', 'S2'],
        'cnt' : [3, 2, 5, 8, 10, 1, 2, 2, 7],
        'val' : ['a', 'n', 'cb', 'mk', 'bg', 'dgb', 'rd', 'cb', 'uyi'],
    })

# I'm not sure if we can rely on val to be a col of unique values so I am just going to 
# make a new column that is the id for this, now on a very large dataframe that wouldn't 
# fit in memory this may not be a reasonable method of creating a unique new column but 
# for the purposes of this example this will be sufficient
df['id'] = np.arange(len(df))

# now let's convert this dataframe into a Dask dataframe
# we will only use 1 partition because this is a small sample and would use more in a real world case
ddf = dd.from_pandas(df, npartitions=1)

# create a function that take the results of the grouped by sub dataframes and returns the row
# where the cnt is greatest
def select_max(grouped_df):
    row_with_max_cnt_index = grouped_df['cnt'].argmax()
    row_with_max_cnt = grouped_df.loc[row_with_max_cnt_index]
    return row_with_max_cnt['id']

# now chain that function into an apply run on the output of the groupby operation
# note: this also may not be the best strategy if the resulting list is too long
# if that is the case, will need to better thread the output of this into the next step
keep_ids = ddf.groupby(['foo', 'bar']).apply(select_max, meta=pd.Series()).compute()

# this is pretty straightforward, just get the rows that match the ids from the max cnt applied method
subset_df = ddf[ddf['id'].isin(keep_ids)]
print(subset_df.compute())
点赞