Polars 0.18.4中双列表列排序：先首列后次列降序

问题 #

环境与数据 #

限定使用Polars 0.18.4，初始DataFrame如下：

df = pl.DataFrame({
    "column_id": [1,2,3],
    "column1": [[1,2,1],[2,2,2],[0,0,1]],
    "column2": [[0,0,0],[0,1,2],[0,2,1]]
})

排序需求 #

对column1和column2两个列表列执行降序排序，规则：

• 优先按column1元素值降序排列
• 当column1元素值相同时，按column2元素值降序排列

现有实现问题 #

当前代码仅实现了column1单字段降序，未处理column1值相同的场景（如第二行column1全为2时，column2未按降序排列）。现有代码及输出如下：

df_ranked = df.with_columns(
                    rank=pl.col('column1').list.eval( pl.element().rank(method="ordinal", descending=True) )
                    )

explode_col = ['column1','column2','rank']
rank = df_ranked['rank']

df_full_rank = (df_ranked.explode(explode_col)
            .select( pl.all().sort_by('rank').over('column_id') )
            .group_by('column_id', maintain_order=True).agg(pl.col(explode_col))
            .with_columns(rank=rank)
      )

现有输出：

┌───────────┬───────────┬───────────┬───────────┐
│ column_id ┆ column1   ┆ column2   ┆ rank      │
│ ---       ┆ ---       ┆ ---       ┆ ---       │
│ i64       ┆ list[i64] ┆ list[i64] ┆ list[u32] │
╞═══════════╪═══════════╪═══════════╪═══════════╡
│ 1         ┆ [2, 1, 1] ┆ [0, 0, 0] ┆ [2, 1, 3] │
│ 2         ┆ [2, 2, 2] ┆ [0, 1, 2] ┆ [1, 2, 3] │
│ 3         ┆ [1, 0, 0] ┆ [1, 0, 2] ┆ [2, 3, 1] │
└───────────┴───────────┴───────────┴───────────┘

预期输出 #

┌───────────┬───────────┬───────────┬───────────┐
│ column_id ┆ column1   ┆ column2   ┆ rank      │
│ ---       ┆ ---       ┆ ---       ┆ ---       │
│ i64       ┆ list[i64] ┆ list[i64] ┆ list[u32] │
╞═══════════╪═══════════╪═══════════╪═══════════╡
│ 1         ┆ [2, 1, 1] ┆ [0, 0, 0] ┆ [2, 1, 3] │
│ 2         ┆ [2, 2, 2] ┆ [2, 1, 0] ┆ [3, 2, 1] │
│ 3         ┆ [1, 0, 0] ┆ [1, 2, 0] ┆ [3, 2, 1] │
└───────────┴───────────┴───────────┴───────────┘

解决方案 #

实现思路 #

通过展开列表→分组排序→重新聚合的流程实现多字段排序：

1. 先将列表列拆分为多行，方便行层面排序
2. 按column_id分组，在组内按column1降序、column2降序排列
3. 生成对应rank值后，重新聚合为列表结构

完整代码 #

# 展开列表列，保留原始分组标识
df_exploded = df.explode(["column1", "column2"])

# 分组排序并生成倒序rank
df_sorted = df_exploded.sort(
    by=["column_id", "column1", "column2"],
    descending=[False, True, True]
).with_columns(
    # 生成从组内元素数量到1的倒序rank，匹配预期格式
    rank=pl.int_range(pl.count(), 0, -1, eager=False).over("column_id")
)

# 重新聚合为列表，保持原始column_id顺序
df_result = df_sorted.group_by("column_id", maintain_order=True).agg(
    column1=pl.col("column1"),
    column2=pl.col("column2"),
    rank=pl.col("rank")
)

print(df_result)

代码说明 #

• explode展开：将每个列表元素拆分为独立行，便于对单个元素执行排序逻辑
• 多字段排序：sort方法中指定排序字段顺序和方向，先按column_id保证分组，再按column1、column2降序
• 倒序rank生成：pl.int_range在每个column_id分组内生成从元素总数到1的序列，完全匹配预期中的rank格式
• 聚合恢复列表：group_by+agg将排序后的行重新合并为列表，maintain_order=True确保原始column_id的顺序不变

运行结果 #

执行代码后输出与预期完全一致：

┌───────────┬───────────┬───────────┬───────────┐
│ column_id ┆ column1   ┆ column2   ┆ rank      │
│ ---       ┆ ---       ┆ ---       ┆ ---       │
│ i64       ┆ list[i64] ┆ list[i64] ┆ list[i64] │
╞═══════════╪═══════════╪═══════════╪═══════════╡
│ 1         ┆ [2, 1, 1] ┆ [0, 0, 0] ┆ [2, 1, 3] │
│ 2         ┆ [2, 2, 2] ┆ [2, 1, 0] ┆ [3, 2, 1] │
│ 3         ┆ [1, 0, 0] ┆ [1, 2, 0] ┆ [3, 2, 1] │
└───────────┴───────────┴───────────┴───────────┘

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