Python Polars分组聚合后如何获取最值对应的时间戳

问题背景 #

现有如下Polars数据集：

df = pl.from_repr("""
┌─────────────────────┬────────────┬────────┐
│ Timestamps          ┆ Date       ┆ Values │
│ ---                 ┆ ---        ┆ ---    │
│ datetime[μs]        ┆ date       ┆ i64    │
╞═════════════════════╪════════════╪════════╡
│ 2022-10-01 04:00:00 ┆ 2022-10-01 ┆ 1      │
│ 2022-10-01 06:00:00 ┆ 2022-10-01 ┆ 2      │
│ 2022-10-01 08:30:00 ┆ 2022-10-01 ┆ 2      │
│ 2022-12-02 09:00:00 ┆ 2022-12-02 ┆ 5      │
│ 2022-12-02 09:15:00 ┆ 2022-12-02 ┆ 4      │
│ 2022-12-02 09:30:00 ┆ 2022-12-02 ┆ 3      │
│ 2022-12-02 09:45:00 ┆ 2022-12-02 ┆ 2      │
│ 2022-12-02 10:00:00 ┆ 2022-12-02 ┆ 1      │
│ 2022-12-02 10:15:00 ┆ 2022-12-02 ┆ 3      │
└─────────────────────┴────────────┴────────┘
""")

需要按Date分组后，获取每组Values的最大值、最小值，以及对应的Timestamps。比如最大值的期望输出为：

┌────────────┬────────┬─────────────────────┐
│ Date       ┆ Values ┆ Timestamps          │
│ ---        ┆ ---    ┆ ---                 │
│ date       ┆ i64    ┆ datetime[μs]        │
╞════════════╪════════╪═════════════════════╡
│ 2022-12-02 ┆ 5      ┆ 2022-12-02 09:00:00 │
│ 2022-10-01 ┆ 2      ┆ 2022-10-01 06:00:00 │
└────────────┴────────┴─────────────────────┘

当前尝试的代码无法提取对应时间戳：

df.group_by('Date').agg(pl.max('Values', 'Timestamps')).sort("Date", descending=True)

解决方案 #

方法1：使用arg_max/arg_min定位行 #

Polars的arg_max和arg_min函数可以返回分组内最值对应的索引，通过这个索引能直接提取对应的Timestamps和Values。

获取最大值及对应时间戳

max_result = (
    df.group_by('Date')
    .agg(
        pl.col('Values').max().alias('max_values'),
        pl.col('Timestamps').arg_max().alias('max_timestamp')
    )
    .sort('Date', descending=True)
)
print(max_result)

输出结果符合预期：

┌────────────┬────────────┬─────────────────────┐
│ Date       ┆ max_values ┆ max_timestamp       │
│ ---        ┆ ---        ┆ ---                 │
│ date       ┆ i64        ┆ datetime[μs]        │
╞════════════╪════════════╪═════════════════════╡
│ 2022-12-02 ┆ 5          ┆ 2022-12-02 09:00:00 │
│ 2022-10-01 ┆ 2          ┆ 2022-10-01 06:00:00 │
└────────────┴────────────┴─────────────────────┘

同时获取最大值、最小值及对应时间戳

如果需要同时获取两组数据，可在agg中添加arg_min相关聚合：

full_result = (
    df.group_by('Date')
    .agg(
        pl.col('Values').max().alias('max_values'),
        pl.col('Timestamps').arg_max().alias('max_timestamp'),
        pl.col('Values').min().alias('min_values'),
        pl.col('Timestamps').arg_min().alias('min_timestamp')
    )
    .sort('Date', descending=True)
)
print(full_result)

输出：

┌────────────┬────────────┬─────────────────────┬────────────┬─────────────────────┐
│ Date       ┆ max_values ┆ max_timestamp       ┆ min_values ┆ min_timestamp       │
│ ---        ┆ ---        ┆ ---                 ┆ ---        ┆ ---                 │
│ date       ┆ i64        ┆ datetime[μs]        ┆ i64        ┆ datetime[μs]        │
╞════════════╪════════════╪═════════════════════╪════════════╪═════════════════════╡
│ 2022-12-02 ┆ 5          ┆ 2022-12-02 09:00:00 ┆ 1          ┆ 2022-12-02 10:00:00 │
│ 2022-10-01 ┆ 2          ┆ 2022-10-01 06:00:00 ┆ 1          ┆ 2022-10-01 04:00:00 │
└────────────┴────────────┴─────────────────────┴────────────┴─────────────────────┘

方法2：使用sort+first/last #

如果分组后先按Values排序，再取首尾行也能得到结果，适合需要保留更多列的场景：

max_result_alt = (
    df.sort('Values', descending=True)
    .group_by('Date')
    .first()
    .select('Date', 'Values', 'Timestamps')
    .sort('Date', descending=True)
)
print(max_result_alt)

这个方法会直接取出每组中Values最大的那一行，输出和期望一致。

说明 #

原代码错误在于pl.max('Values', 'Timestamps')的用法，pl.max只能接收一个列名（或表达式），无法直接关联两个列。需要用arg_max来获取最值对应的行索引，进而提取对应时间戳。

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