Pandas: top N строк, верхних N строк на группу, эквивалентных для ROW_NUMBER OVER (PARTITION BY ... ORDER BY ...)

Question

Что такое эквивалент python функции TOP в T-SQL? Я ищу, чтобы отфильтровать мою фреймворк в верхние 50K строк. Я посмотрел онлайн, и я не могу найти простой пример.

Answer 1

UPDATE: - показывает различные подходы панд, в том числе:

верхних N строк в каждой группу

верхних N строк с офсетной

эквивалента для функций агрегатного SQL:

ROW_NUMBER()/RANK() OVER(PARTITION BY ... ORDER BY ...)

образец DF:

df = pd.DataFrame({ 
    'dep': np.random.choice(list('ABC'), 20), 
    'manager_id': np.random.randint(0, 10, 20), 
    'salary': np.random.randint(5000, 5006, 20) 
}) 

----------------------- Оригинальный DF ------------------- -----

In [2]: df 
Out[2]: 
    dep manager_id salary 
0 B   5 5005 
1 A   6 5001 
2 C   8 5000 
3 A   7 5000 
4 B   0 5002 
5 A   3 5003 
6 A   2 5004 
7 A   2 5004 
8 C   3 5002 
9 C   4 5001 
10 A   9 5002 
11 C   9 5000 
12 B   8 5004 
13 A   1 5003 
14 C   7 5005 
15 B   0 5002 
16 B   2 5003 
17 A   4 5000 
18 B   2 5003 
19 B   7 5003 

------------------ верхние 5 строк (отсортировано по оригинальному указателю) ----------- --------

In [3]: df.head(5) 
Out[3]: 
    dep manager_id salary 
0 B   5 5005 
1 A   6 5001 
2 C   8 5000 
3 A   7 5000 
4 B   0 5002 

--- верхние 5 строк (отсортированы по manager_id DESC, dep ASC) ----

In [4]: df.sort_values(by=['manager_id', 'dep'], ascending=[False,True]).head(5) 
Out[4]: 
    dep manager_id salary 
10 A   9 5002 
11 C   9 5000 
12 B   8 5004 
2 C   8 5000 
3 A   7 5000 

--- эквивалент для SELECT * FROM tab ORDER BY salary DESC LIMIT 5 OFFSET 3 ---

In [19]: df.nlargest(5+3, columns=['salary']).tail(5) 
Out[19]: 
    dep manager_id salary 
7 A   2 5004 
12 B   8 5004 
5 A   3 5003 
13 A   1 5003 
16 B   2 5003 

---- Top 2 зарплаты в каждом отделе (без дубликатов) -----

--- эквивалент для SQL: row_number() over(partition by DEP order by SALARY desc) ---

In [7]: (df.assign(rn=df.sort_values(['salary'], ascending=False) 
    ...:     .groupby(['dep']) 
    ...:     .cumcount() + 1) 
    ...: .query('rn < 3') 
    ...: .sort_values(['dep','rn']) 
    ...:) 
Out[7]: 
    dep manager_id salary rn 
6 A   2 5004 1 
7 A   2 5004 2 
0 B   5 5005 1 
12 B   8 5004 2 
14 C   7 5005 1 
8 C   3 5002 2 

--- Top 2 зарплаты в каждом отделе (с использованием "nlargest") ----

In [15]: df.loc[df.groupby('dep')['salary'].nlargest(2).reset_index()['level_1']] 
Out[15]: 
    dep manager_id salary 
6 A   2 5004 
7 A   2 5004 
0 B   5 5005 
12 B   8 5004 
14 C   7 5005 
8 C   3 5002 

--- второй и третий самые высокие зарплаты в каждом отделе ---

In [16]: (df.assign(rn=df.sort_values(['salary'], ascending=False) 
    ....:     .groupby(['dep']) 
    ....:     .cumcount() + 1) 
    ....: .query('rn >= 2 and rn <= 3') 
    ....: .sort_values(['dep','rn']) 
    ....:) 
Out[16]: 
    dep manager_id salary rn 
7 A   2 5004 2 
13 A   1 5003 3 
12 B   8 5004 2 
18 B   2 5003 3 
8 C   3 5002 2 
9 C   4 5001 3 

--- верхние 2 зарплаты в каждом отделе (с дубликатами) ----

--- эквивалент для SQL : rank() over(partition by DEP order by SALARY desc) ---

In [18]: (df.assign(rnk=df.groupby(['dep'])['salary'] 
    ....:     .rank(method='min', ascending=False)) 
    ....: .query('rnk < 3') 
    ....: .sort_values(['dep','rnk']) 
    ....:) 
Out[18]: 
    dep manager_id salary rnk 
6 A   2 5004 1.0 
7 A   2 5004 1.0 
0 B   5 5005 1.0 
12 B   8 5004 2.0 
14 C   7 5005 1.0 
8 C   3 5002 2.0