如何用Polars列选择器cs.datetime()获取含最大日期值的列名

问题背景 #

需要找出Polars DataFrame每行中包含最大日期值的列名，由于日期列数量较多，手动指定列名过于繁琐，希望通过列选择器cs.datetime()自动筛选日期列来实现。

解决方案 #

问题分析 #

原代码中arg_max_horizontal(pl.select(cs.datetime()))存在错误：pl.select(cs.datetime())返回的是DataFrame对象，而函数需要的是多个列表达式（pl.Expr）作为参数。我们需要先获取所有日期列的表达式，再传入函数，或者改用更简洁的Polars原生方法实现。

方法一：修正原函数调用 #

先通过cs.datetime()获取所有日期列的表达式，再将其拆分为可变参数传入自定义函数：

import polars as pl
from datetime import datetime
import polars.selectors as cs

data = {
    "ID" : [1,2,3],
    "Values_A" : [datetime(1,1,2),datetime(1,1,3),datetime(1,1,4)],
    "Values_B" : [datetime(1,1,4),datetime(1,1,7),datetime(1,1,2)]
}

df = pl.DataFrame(data)

def arg_max_horizontal(*columns: pl.Expr) -> pl.Expr:
    return (
        pl.concat_list(columns)
        .list.arg_max()
        .replace_strict({i: col.meta.output_name() for i, col in enumerate(columns)})
    )

# 获取所有日期列的表达式，拆分为可变参数传入函数
date_cols = cs.datetime()
result = df.with_columns(
    Largest=arg_max_horizontal(*date_cols)
)

print(result)

方法二：使用Polars原生struct简洁实现 #

利用pl.struct将日期列打包，再通过struct.arg_max()直接获取最大值对应的字段名，无需自定义函数：

import polars as pl
from datetime import datetime
import polars.selectors as cs

data = {
    "ID" : [1,2,3],
    "Values_A" : [datetime(1,1,2),datetime(1,1,3),datetime(1,1,4)],
    "Values_B" : [datetime(1,1,4),datetime(1,1,7),datetime(1,1,2)]
}

df = pl.DataFrame(data)

result = df.with_columns(
    Largest=pl.struct(cs.datetime()).arg_max()
)

print(result)

两种方法均会输出如下结果：

shape: (3, 3)
┌─────┬─────────────┬─────────────┬─────────┐
│ ID  ┆ Values_A    ┆ Values_B    ┆ Largest │
│ --- ┆ ---         ┆ ---         ┆ ---     │
│ i64 ┆ datetime[μs]┆ datetime[μs]┆ str     │
╞═════╪═════════════╪═════════════╪═════════╡
│ 1   ┆ 0001-01-02  ┆ 0001-01-04  ┆ Values_B│
│ 2   ┆ 0001-01-03  ┆ 0001-01-07  ┆ Values_B│
│ 3   ┆ 0001-01-04  ┆ 0001-01-02  ┆ Values_A│
└─────┴─────────────┴─────────────┴─────────┘

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