如何基于外部列表对Polars DataFrame进行排序

早上好呀！你说的这个需求完全可以实现哒～我来给你具体讲讲怎么做。

首先先确认下你的示例场景，你的DataFrame和列表是这样的：

import polars as pl

df = pl.from_repr("""
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ str │
╞═════╪═════╡
│ 1   ┆ x   │
│ 2   ┆ y   │
│ 3   ┆ z   │
│ 4   ┆ p   │
│ 5   ┆ f   │
└─────┴─────┘
""")

l = [1,3,5,2,4]

接下来给你两种实用的方法：

• 方法一：字典映射法（简单高效，推荐）
  我们可以先把列表里的每个元素和它的位置索引做个映射，然后给DataFrame加一个临时的排序列，用这个列排序后再删掉它就搞定了：

# 创建元素到索引的映射字典
sort_map = {val: idx for idx, val in enumerate(l)}

# 添加临时排序列、排序、删除临时列
sorted_df = df.with_columns(
    pl.col("a").map_dict(sort_map).alias("sort_order")
).sort("sort_order").drop("sort_order")

print(sorted_df)

运行后就能得到按列表l顺序排列的结果：

┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ str │
╞═════╪═════╡
│ 1   ┆ x   │
│ 3   ┆ z   │
│ 5   ┆ f   │
│ 2   ┆ y   │
│ 4   ┆ p   │
└─────┴─────┘

• 方法二：关联排序法（适合复杂场景）
  如果你的排序逻辑更复杂，比如要基于多列匹配外部顺序，可以试试把列表转成带索引的小DataFrame，再和原表关联后排序：

# 将列表转为带索引的DataFrame
order_df = pl.Series(l).to_frame("a").with_row_index("sort_idx")

# 关联后排序，恢复原列顺序
sorted_df = order_df.join(df, on="a").sort("sort_idx").select(df.columns)

print(sorted_df)

这个方法的结果和上面完全一样，灵活性更强哦～

备注：内容来源于stack exchange，提问作者Ghost