基于整洁Polars DataFrame计算横截面排名及优化问询

问题1：将ranks逆透视为整洁格式DataFrame #

你可以通过两次逆透视（melt）+ 合并的方式实现，先处理收益列，再拆分排名列表并处理，最后将两者按日期和标的关联：

# 第一步：将宽表的return列逆透视成长表
melted_returns = ranks.melt(
    id_vars="date",
    value_name="return",
    variable_name="symbol"
)

# 第二步：把rank列表转换为结构体后展开，再逆透视
rank_struct = ranks.select(
    "date",
    pl.col("rank").list.to_struct(fields=ranks.columns[1:-1])
).unnest("rank")

melted_ranks = rank_struct.melt(
    id_vars="date",
    value_name="rank",
    variable_name="symbol"
)

# 第三步：合并收益和排名数据
final_df = melted_returns.join(melted_ranks, on=["date", "symbol"]).sort(["date", "symbol"])
print(final_df)

输出结果：

shape: (18, 4)
┌────────────┬─────────┬──────────┬──────┐
│ date       ┆ symbol  ┆ return   ┆ rank │
│ ---        ┆ ---     ┆ ---      ┆ ---  │
│ str        ┆ str     ┆ f64      ┆ u8   │
╞════════════╪═════════╪══════════╪══════╡
│ 2023-12-30 ┆ symbol1 ┆ null     ┆ null │
│ 2023-12-30 ┆ symbol2 ┆ null     ┆ null │
│ 2023-12-30 ┆ symbol3 ┆ null     ┆ null │
│ 2023-12-31 ┆ symbol1 ┆ null     ┆ null │
│ 2023-12-31 ┆ symbol2 ┆ null     ┆ null │
│ …          ┆ …       ┆ …        ┆ …    │
│ 2024-01-05 ┆ symbol2 ┆ 0.095238 ┆ 1    │
│ 2024-01-05 ┆ symbol3 ┆ 0.0625   ┆ 3    │
│ 2024-01-06 ┆ symbol1 ┆ 0.086957 ┆ 2    │
│ 2024-01-06 ┆ symbol2 ┆ 0.090909 ┆ 1    │
│ 2024-01-06 ┆ symbol3 ┆ null     ┆ null │
└────────────┴─────────┴──────────┴──────┘

问题2：无需透视，直接在长格式DataFrame上高效计算 #

完全不需要先透视，Polars支持直接按date分组计算横截面排名，这是更高效的方案（尤其适合百万级数据）：

result = df.with_columns(
    # 先计算每个标的的周期收益
    pl.col("price").pct_change(n=2).over("symbol").alias("return")
).with_columns(
    # 按日期分组，对收益降序排名（自动忽略null值）
    pl.col("return")
    .rank(descending=True, method="ordinal")
    .over("date")
    .cast(pl.UInt8)
    .alias("rank")
).sort(["date", "symbol"])

print(result)

输出结果（和问题1的最终结果逻辑一致，且保留原数据的所有行）：

shape: (16, 5)
┌─────────┬────────────┬───────┬──────────┬──────┐
│ symbol  ┆ date       ┆ price ┆ return   ┆ rank │
│ ---     ┆ ---        ┆ ---   ┆ ---      ┆ ---  │
│ str     ┆ str        ┆ i64   ┆ f64      ┆ u8   │
╞═════════╪════════════╪═══════╪══════════╪══════╡
│ symbol1 ┆ 2023-12-30 ┆ 100   ┆ null     ┆ null │
│ symbol2 ┆ 2023-12-30 ┆ 200   ┆ null     ┆ null │
│ symbol3 ┆ 2023-12-30 ┆ 3000  ┆ null     ┆ null │
│ symbol1 ┆ 2023-12-31 ┆ 105   ┆ null     ┆ null │
│ symbol3 ┆ 2023-12-31 ┆ 3100  ┆ null     ┆ null │
│ …       ┆ …          ┆ …     ┆ …        ┆ …    │
│ symbol1 ┆ 2024-01-05 ┆ 120   ┆ 0.090909 ┆ 2    │
│ symbol2 ┆ 2024-01-05 ┆ 230   ┆ 0.095238 ┆ 1    │
│ symbol3 ┆ 2024-01-05 ┆ 3400  ┆ 0.0625   ┆ 3    │
│ symbol1 ┆ 2024-01-06 ┆ 125   ┆ 0.086957 ┆ 2    │
│ symbol2 ┆ 2024-01-06 ┆ 240   ┆ 0.090909 ┆ 1    │
└─────────┴────────────┴───────┴──────────┴──────┘

性能说明 #

这种直接在长表上用over("date")分组排名的方式，避免了透视/逆透视带来的数据结构转换开销，Polars的矢量化分组操作可以高效处理百万级行数据，远优于先透视再处理的方案。

内容的提问来源于stack exchange，提问作者Andi