在Python Pandas数据框中使用类似SQL的窗口函数进行行编号

12 浏览
0 Comments

在Python Pandas数据框中使用类似SQL的窗口函数进行行编号

我来自于SQL背景，我经常使用以下数据处理步骤：

  1. 根据一个或多个字段对数据表进行分区
  2. 对每个分区，为其中的每一行添加一个行号，根据一个或多个其他字段进行排序，分析师可以指定升序或降序

例如：

df = pd.DataFrame({'key1' : ['a','a','a','b','a'],
           'data1' : [1,2,2,3,3],
           'data2' : [1,10,2,3,30]})
df
     data1        data2     key1    
0    1            1         a           
1    2            10        a        
2    2            2         a       
3    3            3         b       
4    3            30        a

我正在寻找如何在Pandas中实现类似于SQL窗口函数的功能：

RN = ROW_NUMBER() OVER (PARTITION BY Key1 ORDER BY Data1 ASC, Data2 DESC)
    data1        data2     key1    RN
0    1            1         a       1    
1    2            10        a       2 
2    2            2         a       3
3    3            3         b       1
4    3            30        a       4

我尝试了以下方法，在没有“分区”的情况下可以工作：

def row_number(frame,orderby_columns, orderby_direction,name):
    frame.sort_index(by = orderby_columns, ascending = orderby_direction, inplace = True)
    frame[name] = list(xrange(len(frame.index)))

我试图扩展这个思路以适应分区（在pandas中是分组），但是以下方法没有成功：

df1 = df.groupby('key1').apply(lambda t: t.sort_index(by=['data1', 'data2'], ascending=[True, False], inplace = True)).reset_index()
def nf(x):
    x['rn'] = list(xrange(len(x.index)))
df1['rn1'] = df1.groupby('key1').apply(nf)

但是当我这样做时，只得到了很多NaN值。

理想情况下，有一种简洁的方法可以复制SQL中的窗口函数功能（我已经想出了基于窗口的聚合函数...在pandas中只需要一行代码就可以实现）...有人可以告诉我在Pandas中实现这样编号行的最惯用的方法吗？

0