Python实现窗口函数

计算每只股票持仓前五的用户信息

SQL实现

select * from 
(
select sec_code,acct_id,hold_vol,hold_amt,DENSE_RANK() OVER(partition by sec_code order by hold_amt desc) as rank 
from sec_hold 
) a
where a.rank <=5

输出

#Python实现

def query_sql(sql):
    conn = psycopg2.connect(database='quant', host='10.211.55.20', port='5432', user='kiwi',password='yx3239736')
    conn.set_client_encoding('utf-8')
    cursor = conn.cursor()
    cursor.execute(sql)
    rows = cursor.fetchall()
    df = pd.DataFrame(rows)
    return df

def group_func(dataFame):
    return dataFame.sort_values(['hold_amt'],ascending=False)[:5]


sql = "select sec_code,acct_id,hold_vol,hold_amt from sec_hold "
df = query_sql(sql)
df.rename(columns={0:'sec_code', 1:'acct_id', 2:'hold_vol',3:'hold_amt',4:'rank'}, inplace = True)

df1 = df.groupby('sec_code').apply(group_func)

print(df1.head(30))

输出

结果一致