如何替代map_groups实现Polars数据框的双层排序需求？

需求说明 #

需要对Polars DataFrame完成两步排序逻辑：

1. 整体按会话的最新时间戳降序排列，让最新产生的会话组排在顶部
2. 每个会话组内部按时间戳升序排列，还原用户操作时序

原始DataFrame结构：

session_id | timestamp
abcd       | 2024-04-01 
abcd       | 2024-04-06
abcd       | 2024-04-03
efgh       | 2024-04-05
efgh       | 2024-04-04
efgh       | 2024-04-07

以下是两种无需使用map_groups的实现方案：

方案1：窗口函数+复合排序 #

通过窗口函数标记每个会话的最新时间，再通过复合排序规则一次性完成需求：

import polars as pl

# 构造原始数据
df = pl.DataFrame({
    "session_id": ["abcd", "abcd", "abcd", "efgh", "efgh", "efgh"],
    "timestamp": ["2024-04-01", "2024-04-06", "2024-04-03", "2024-04-05", "2024-04-04", "2024-04-07"]
}).with_columns(pl.col("timestamp").str.to_date())

# 实现排序逻辑
result = df.with_columns(
    # 标记每个会话的最新时间戳
    pl.col("timestamp").max().over("session_id").alias("max_session_ts")
).sort(
    # 排序规则：先按会话最新时间降序，再按会话ID，最后按组内时间升序
    by=["max_session_ts", "session_id", "timestamp"],
    descending=[True, False, False]
).drop("max_session_ts")  # 删除辅助列

print(result)

输出结果：

session_id | timestamp
efgh       | 2024-04-04
efgh       | 2024-04-05
efgh       | 2024-04-07
abcd       | 2024-04-01 
abcd       | 2024-04-03
abcd       | 2024-04-06

方案2：自定义会话顺序+分类排序 #

先确定会话组的排序优先级，再将会话ID转为自定义顺序的分类类型，实现整体排序：

import polars as pl

# 构造原始数据（同上）
df = pl.DataFrame({
    "session_id": ["abcd", "abcd", "abcd", "efgh", "efgh", "efgh"],
    "timestamp": ["2024-04-01", "2024-04-06", "2024-04-03", "2024-04-05", "2024-04-04", "2024-04-07"]
}).with_columns(pl.col("timestamp").str.to_date())

# 第一步：获取会话组的排序顺序（按会话最新时间降序）
session_order = df.group_by("session_id").agg(
    pl.col("timestamp").max().alias("max_ts")
).sort("max_ts", descending=True).select("session_id")

# 第二步：按自定义会话顺序+组内时间升序排序
result = df.sort(
    by=[
        # 将session_id转为自定义顺序的分类类型
        pl.col("session_id").cast(pl.Categorical).cat.set_ordering("custom").cat.set_categories(session_order["session_id"]),
        "timestamp"
    ],
    descending=[False, False]
)

print(result)

该方案输出结果与方案1完全一致，适合需要复用会话排序顺序的场景。

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