【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)