如何使用Polars的rank列对list类型列进行自定义排序？

问题背景 #

现有如下Polars DataFrame：

import polars as pl

data = {
    "user": ["xxx", "yyy", "zzz"],
    "list_1": [[1, 3, 5], [4, 9, 5], [4, 6, 1]],
    "list_2": [[8, 3, 5], [3, 4, 5], [9, 3, 6]],
    "list_3": [[6, 7, 8], [7, 8, 3], [4, 3, 2]],
    "rank": [[1, 2, 3], [1, 3, 2], [2, 3, 1]]
}

df = pl.DataFrame(data)

原始DataFrame输出：

shape: (3, 5)
┌──────┬───────────┬───────────┬───────────┬───────────┐
│ user ┆ list_1    ┆ list_2    ┆ list_3    ┆ rank      │
│ ---  ┆ ---       ┆ ---       ┆ ---       ┆ ---       │
│ str  ┆ list[i64] ┆ list[i64] ┆ list[i64] ┆ list[i64] │
╞══════╪═══════════╪═══════════╪═══════════╪═══════════╡
│ xxx  ┆ [1, 3, 5] ┆ [8, 3, 5] ┆ [6, 7, 8] ┆ [1, 2, 3] │
│ yyy  ┆ [4, 9, 5] ┆ [3, 4, 5] ┆ [7, 8, 3] ┆ [1, 3, 2] │
│ zzz  ┆ [4, 6, 1] ┆ [9, 3, 6] ┆ [4, 3, 2] ┆ [2, 3, 1] │
└──────┴───────────┴───────────┴───────────┴───────────┘

其中rank列由list_1的排名生成，需要使用rank列对list_1、list_2、list_3按行排序，期望输出如下：

shape: (3, 5)
┌──────┬───────────┬───────────┬───────────┬───────────┐
│ user ┆ list_1    ┆ list_2    ┆ list_3    ┆ rank      │
│ ---  ┆ ---       ┆ ---       ┆ ---       ┆ ---       │
│ str  ┆ list[i64] ┆ list[i64] ┆ list[i64] ┆ list[i64] │
╞══════╪═══════════╪═══════════╪═══════════╪═══════════╡
│ xxx  ┆ [1, 3, 5] ┆ [8, 3, 5] ┆ [6, 7, 8] ┆ [1, 2, 3] │
│ yyy  ┆ [4, 5, 9] ┆ [3, 5, 4] ┆ [7, 3, 8] ┆ [1, 3, 2] │
│ zzz  ┆ [1, 4, 6] ┆ [6, 9, 3] ┆ [2, 4, 3] ┆ [2, 3, 1] │
└──────┴───────────┴───────────┴───────────┴───────────┘

解决方案 #

可以利用Polars的list.gather方法，结合list.eval将1-based的rank转为0-based索引，实现按行排序列表列：

df = df.with_columns(
    # 对每个列表列，用rank转成的索引重新排列元素
    pl.col(col).list.gather(pl.col("rank").list.eval(pl.element() - 1)).alias(col)
    for col in ["list_1", "list_2", "list_3"]
)

print(df)

代码解释 #

1. 索引转换：pl.col("rank").list.eval(pl.element() - 1) 将rank列的1-based排名值转为0-based索引，因为Polars列表的索引从0开始。
2. 按索引重排：list.gather方法接收索引列表，将原列表中的元素按索引顺序重新排列，实现按rank指定的顺序排序。
3. 批量处理：通过生成器表达式批量处理list_1、list_2、list_3三个列，简化代码。

运行上述代码后，输出结果将与期望一致。

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