【pandas】[4] DataFrame实现sql中row_number() over(partition by column_1 order by column_2)

需求：pandas中能不能实现如sql中一样的分组排序取值

 

1、构建测试数据

import pandas as pd
df = pd.DataFrame({'k1' : ['a1','a2','a1','b1','b2'],
    'k2' : ['c1','d1','c1','c2','d2'],
    'data' : [10,100,20,30,300]})
print(df)

   k1  k2  data
0  a1  c1    10
1  a2  d1   100
2  a1  c1    20
3  b1  c2    30
4  b2  d2   300

2、使用pandas中的groupby与rank函数实现sql中row_number() over()的功能

df['row_number'] = df.groupby(['k1'])['data'].rank(ascending=False,method='dense')

print(df)

   k1  k2  data  row_number
0  a1  c1    10         2.0
1  a2  d1   100         1.0
2  a1  c1    20         1.0
3  b1  c2    30         1.0
4  b2  d2   300         1.0

可以看出如上df中的k1列的第一行和第三行的'a1'取值是重复的。通过第2步的实现，新增一列row_number列。

至此，完成通过pandas来实现sql中的row_number() over()的功能

 

封装成函数

def row_number(df, groupby_col=[], orderby_col='', ascending=True):
    '''

    :param df: 需要处理的数据集；pandas.DataFrame
    :param groupby_col: 需要分组的列；list
    :param orderby_col: 需要分组后，进行排序的列；columns_names
    :param ascending: 排序方式，默认升序
    :return: pandas.Series
    '''

    return df.groupby(groupby_col)[orderby_col].rank(ascending=False,method='dense')

df['row_numbers'] = row_number(df, groupby_col=['k1'], orderby_col='data', ascending=False)

print(df)