如何在Polars中按指定规则排序行：货币降序+自定义字母顺序

问题场景 #

现有如下Polars数据框：

import polars as pl

df = pl.DataFrame({
    "currency": ["EUR","EUR","EUR","USD","USD","USD"],
    "alphabet": ["A","B","C","A","B","C"]
})

需要按两个规则排序：

1. currency列降序排列（USD优先于EUR）
2. alphabet列按C < A < B的自定义顺序排列

解决方案 #

方法一：利用分类类型（Categorical）实现自定义排序

这是Polars推荐的方式，通过将列转换为分类类型并指定自定义顺序，直接利用内置排序逻辑：

import polars as pl

df = pl.DataFrame({
    "currency": ["EUR","EUR","EUR","USD","USD","USD"],
    "alphabet": ["A","B","C","A","B","C"]
})

# 定义alphabet的自定义排序顺序
custom_alphabet_order = ["C", "A", "B"]

# 将alphabet转为分类类型，设置按物理顺序排序
sorted_df = df.with_columns(
    pl.col("alphabet")
    .cast(pl.Categorical)
    .cat.set_ordering("physical")  # 启用物理顺序（即我们定义的顺序）
).sort(
    by=["currency", "alphabet"],
    descending=[True, False]  # currency降序，alphabet按自定义顺序升序
)

print(sorted_df)

方法二：通过映射权重实现自定义排序

如果不想修改列类型，可以通过映射字典生成排序权重，间接实现自定义顺序：

import polars as pl

df = pl.DataFrame({
    "currency": ["EUR","EUR","EUR","USD","USD","USD"],
    "alphabet": ["A","B","C","A","B","C"]
})

# 定义映射字典：值越小排序越靠前
order_weight = {"C": 0, "A": 1, "B": 2}

sorted_df = df.with_columns(
    pl.col("alphabet").map_dict(order_weight).alias("sort_key")  # 生成临时排序键
).sort(
    by=["currency", "sort_key"],
    descending=[True, False]
).drop("sort_key")  # 移除临时列

print(sorted_df)

预期输出 #

两种方法都能得到如下结果：

shape: (6, 2)
┌──────────┬──────────┐
│ currency ┆ alphabet │
│ ---      ┆ ---      │
│ str      ┆ str      │
╞══════════╪══════════╡
│ USD      ┆ C        │
│ USD      ┆ A        │
│ USD      ┆ B        │
│ EUR      ┆ C        │
│ EUR      ┆ A        │
│ EUR      ┆ B        │
└──────────┴──────────┘

内容的提问来源于stack exchange，提问作者Muhammad D Vikar